From 187144d4e74a0cb54c06e1aba11b6a3507195997 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 8 Jun 2026 23:09:51 -0700 Subject: [PATCH] Reduce and tighten code comments and docstrings repo-wide (#6095) Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison. --- scripts/check_frontend_dep_removal.py | 329 +++----- scripts/check_new_install_scripts.py | 77 +- scripts/enforce_kwargs_spacing.py | 116 +-- scripts/lint_workflow_triggers.py | 40 +- scripts/lockfile_supply_chain_audit.py | 211 ++--- scripts/notebook_to_python.py | 69 +- scripts/notebook_validator.py | 138 ++-- scripts/run_ruff_format.py | 6 +- scripts/scan_npm_packages.py | 215 ++---- scripts/scan_packages.py | 261 ++----- scripts/stamp_studio_release.py | 7 +- scripts/verify_comment_only_diff.py | 13 +- scripts/verify_import_hoist.py | 64 +- studio/backend/_platform_compat.py | 15 +- studio/backend/auth/authentication.py | 10 +- studio/backend/auth/hashing.py | 2 +- studio/backend/auth/storage.py | 51 +- studio/backend/colab.py | 42 +- studio/backend/core/__init__.py | 4 +- studio/backend/core/_torchao_stub.py | 43 +- studio/backend/core/data_recipe/__init__.py | 4 +- .../backend/core/data_recipe/huggingface.py | 3 +- .../backend/core/data_recipe/jobs/manager.py | 14 +- studio/backend/core/data_recipe/jobs/parse.py | 2 +- studio/backend/core/data_recipe/jobs/types.py | 5 +- .../backend/core/data_recipe/jobs/worker.py | 6 +- .../data_recipe/local_callable_validators.py | 2 +- studio/backend/core/data_recipe/service.py | 11 +- studio/backend/core/export/__init__.py | 9 +- studio/backend/core/export/export.py | 58 +- studio/backend/core/export/orchestrator.py | 76 +- studio/backend/core/export/worker.py | 101 +-- studio/backend/core/inference/_html_to_md.py | 56 +- .../core/inference/anthropic_compat.py | 14 +- studio/backend/core/inference/audio_codecs.py | 38 +- studio/backend/core/inference/defaults.py | 2 +- .../core/inference/external_provider.py | 722 ++++++------------ studio/backend/core/inference/inference.py | 150 ++-- studio/backend/core/inference/key_exchange.py | 9 +- studio/backend/core/inference/llama_cpp.py | 627 ++++++--------- .../core/inference/llama_server_args.py | 26 +- studio/backend/core/inference/mcp_client.py | 41 +- .../backend/core/inference/mlx_inference.py | 56 +- studio/backend/core/inference/orchestrator.py | 112 ++- studio/backend/core/inference/pricing.py | 12 +- studio/backend/core/inference/providers.py | 116 +-- .../core/inference/safetensors_agentic.py | 16 +- .../core/inference/tool_call_parser.py | 22 +- studio/backend/core/inference/tools.py | 239 +++--- studio/backend/core/inference/worker.py | 85 +-- studio/backend/core/tool_healing.py | 53 +- studio/backend/core/training/__init__.py | 2 +- studio/backend/core/training/trainer.py | 212 +++-- studio/backend/core/training/training.py | 65 +- studio/backend/core/training/worker.py | 223 ++---- studio/backend/loggers/config.py | 17 +- studio/backend/loggers/handlers.py | 20 +- studio/backend/main.py | 120 ++- studio/backend/models/inference.py | 143 ++-- studio/backend/models/training.py | 11 +- .../__init__.py | 5 +- .../data_designer_github_repo_seed/impl.py | 14 +- .../data_designer_github_repo_seed/scraper.py | 4 +- .../scraper_impl/gh_client.py | 9 +- .../scraper_impl/scraper.py | 18 +- .../chunking.py | 2 +- studio/backend/routes/auth.py | 19 +- studio/backend/routes/chat_history.py | 19 +- studio/backend/routes/data_recipe/jobs.py | 139 ++-- studio/backend/routes/data_recipe/seed.py | 3 +- studio/backend/routes/data_recipe/validate.py | 15 +- studio/backend/routes/datasets.py | 67 +- studio/backend/routes/export.py | 49 +- studio/backend/routes/inference.py | 2 - studio/backend/routes/mcp_servers.py | 9 +- studio/backend/routes/models.py | 306 +++----- studio/backend/routes/providers.py | 29 +- studio/backend/routes/training.py | 30 +- studio/backend/run.py | 149 ++-- studio/backend/startup_banner.py | 2 +- studio/backend/storage/mcp_servers_db.py | 2 +- studio/backend/storage/providers_db.py | 11 +- studio/backend/storage/studio_db.py | 58 +- studio/backend/tests/conftest.py | 55 +- .../backend/tests/test_anthropic_cache_ttl.py | 21 +- .../backend/tests/test_anthropic_citations.py | 3 +- .../tests/test_anthropic_code_execution.py | 39 +- .../tests/test_anthropic_compaction.py | 91 +-- .../backend/tests/test_anthropic_messages.py | 5 +- .../test_anthropic_thinking_translation.py | 20 +- .../tests/test_anthropic_tool_versions.py | 37 +- .../backend/tests/test_anthropic_web_fetch.py | 18 +- .../tests/test_cache_case_resolution.py | 3 +- .../backend/tests/test_cached_gguf_routes.py | 30 +- .../tests/test_chat_history_storage.py | 3 +- .../test_cleanup_cancelled_checkpoints.py | 9 +- studio/backend/tests/test_desktop_auth.py | 8 +- .../backend/tests/test_export_log_cursor.py | 58 +- .../test_external_provider_usage_chunk.py | 25 +- .../backend/tests/test_frontend_resolution.py | 4 +- .../tests/test_gguf_reload_inheritance.py | 7 +- studio/backend/tests/test_gguf_routing.py | 7 +- studio/backend/tests/test_host_defaults.py | 6 +- .../tests/test_index_bootstrap_origin.py | 16 +- .../tests/test_inference_model_validation.py | 2 +- .../backend/tests/test_kv_cache_estimation.py | 39 +- .../test_lemonade_llamacpp_rocm_bins_mock.py | 10 +- .../test_llama_cpp_cache_aware_disk_check.py | 36 +- .../tests/test_llama_cpp_context_fit.py | 45 +- .../backend/tests/test_llama_cpp_freshness.py | 7 +- .../tests/test_llama_cpp_load_progress.py | 37 +- .../test_llama_cpp_load_progress_live.py | 21 +- .../test_llama_cpp_load_progress_matrix.py | 37 +- .../test_llama_cpp_max_context_threshold.py | 10 - .../tests/test_llama_cpp_mtp_detection.py | 8 +- .../tests/test_llama_cpp_no_context_shift.py | 5 +- .../tests/test_llama_cpp_wait_for_health.py | 11 +- .../test_llama_cpp_wait_for_vram_settle.py | 12 +- .../test_llama_cpp_windows_nvidia_path.py | 29 +- .../backend/tests/test_llama_server_args.py | 13 +- studio/backend/tests/test_login_rate_limit.py | 14 +- studio/backend/tests/test_mcp_servers.py | 44 +- .../tests/test_mcp_stdio_improvements.py | 3 +- studio/backend/tests/test_mcp_stdio_pr5863.py | 13 +- studio/backend/tests/test_middleware.py | 6 - .../backend/tests/test_multimodal_document.py | 43 +- .../tests/test_offline_gguf_cache_fallback.py | 56 +- .../tests/test_openai_citation_markers.py | 17 +- .../tests/test_openai_code_execution.py | 44 +- .../tests/test_openai_image_generation.py | 16 +- .../test_openai_responses_translation.py | 2 +- .../tests/test_openai_tool_passthrough.py | 48 +- studio/backend/tests/test_pricing_edge.py | 1 - studio/backend/tests/test_providers_api.py | 50 +- studio/backend/tests/test_pytorch_mirror.py | 4 +- .../test_recommended_folders_permission.py | 3 +- studio/backend/tests/test_responses_api.py | 12 +- studio/backend/tests/test_rocm_oom_guard.py | 14 +- .../test_safetensors_capability_advertise.py | 20 +- .../tests/test_safetensors_tool_loop.py | 76 +- studio/backend/tests/test_sandbox_tools.py | 35 +- studio/backend/tests/test_studio_api.py | 42 +- .../tests/test_studio_train_validation.py | 6 +- .../tests/test_training_worker_flash_attn.py | 89 +-- .../tests/test_transformers_version.py | 6 +- studio/backend/tests/test_utils.py | 21 +- studio/backend/tests/test_vision_cache.py | 34 +- .../tests/test_windows_gpu_detection_mock.py | 20 +- studio/backend/utils/cache_cleanup.py | 15 +- studio/backend/utils/cpu_threads.py | 4 +- studio/backend/utils/datasets/__init__.py | 22 +- .../backend/utils/datasets/chat_templates.py | 24 +- .../backend/utils/datasets/data_collators.py | 40 +- .../utils/datasets/dataset_none_detect.py | 8 +- .../backend/utils/datasets/dataset_utils.py | 7 +- .../utils/datasets/format_conversion.py | 19 +- .../utils/datasets/format_detection.py | 152 ++-- studio/backend/utils/datasets/llm_assist.py | 140 +--- .../backend/utils/datasets/model_mappings.py | 2 +- studio/backend/utils/datasets/raw_text.py | 4 +- .../backend/utils/datasets/vlm_processing.py | 6 +- studio/backend/utils/hardware/__init__.py | 4 +- studio/backend/utils/hardware/amd.py | 86 +-- studio/backend/utils/hardware/hardware.py | 171 ++--- studio/backend/utils/hardware/nvidia.py | 9 +- .../backend/utils/hardware/vram_estimation.py | 91 +-- .../utils/inference/inference_config.py | 52 +- studio/backend/utils/llama_cpp_freshness.py | 5 +- studio/backend/utils/models/__init__.py | 4 +- studio/backend/utils/models/checkpoints.py | 21 +- studio/backend/utils/models/model_config.py | 314 +++----- studio/backend/utils/paths/__init__.py | 4 +- studio/backend/utils/paths/path_utils.py | 21 +- studio/backend/utils/paths/storage_roots.py | 20 +- studio/backend/utils/subprocess_compat.py | 5 +- studio/backend/utils/transformers_version.py | 36 +- studio/backend/utils/utils.py | 23 +- studio/backend/utils/wheel_utils.py | 14 +- studio/install_llama_prebuilt.py | 18 +- studio/install_python_stack.py | 20 +- tests/_zoo_aggressive_cuda_spoof.py | 57 +- tests/conftest.py | 13 +- tests/python/test_cross_platform_parity.py | 3 +- tests/python/test_e2e_no_torch_sandbox.py | 26 +- ...sentence_transformer_redirect_lifecycle.py | 4 +- tests/python/test_install_python_stack.py | 5 +- tests/python/test_studio_import_no_torch.py | 1 - .../test_tokenizers_and_torch_constraint.py | 7 +- .../saving/gpt-oss-merge/test_merged_model.py | 2 +- tests/saving/gpt-oss-merge/train_and_merge.py | 3 +- .../test_merge_4bit_validation.py | 5 +- .../test_merge_model_perplexity_llama-3.2.py | 20 +- .../test_merge_model_perplexity_mistral.py | 28 +- .../test_merge_model_perplexity_phi_4.py | 23 +- ...st_merged_model_perplexity_llama-3.1-8b.py | 20 +- .../test_merged_model_perplexity_qwen_2.5.py | 25 +- .../test_push_to_hub_merged.py | 3 +- ...t_push_to_hub_merged_sharded_index_file.py | 1 - .../test_save_merged_grpo_model.py | 37 +- .../saving/non_peft/test_mistral_non_peft.py | 5 +- .../saving/non_peft/test_whisper_non_peft.py | 3 - tests/saving/test_unsloth_save.py | 39 +- .../saving/text_to_speech_models/test_csm.py | 6 +- .../text_to_speech_models/test_orpheus.py | 4 +- .../text_to_speech_models/test_whisper.py | 4 +- .../test_index_file_sharded_model.py | 25 +- .../vision_models/test_push_to_hub_merged.py | 7 +- ..._merge_qwen2.5vl32B_model_ocr_benchmark.py | 24 +- ...t_save_merge_vision_model_ocr_benchmark.py | 40 +- tests/security/conftest.py | 19 +- tests/security/fixtures/_build.py | 25 +- .../test_lockfile_supply_chain_audit.py | 56 +- tests/security/test_new_install_scripts.py | 8 +- tests/security/test_scan_npm_packages.py | 31 +- tests/security/test_scan_packages.py | 46 +- tests/studio/_playwright_robust.py | 245 ++---- .../smoke_test_parallel_studio_home.py | 53 +- .../install/test_llama_pr_force_and_source.py | 11 +- tests/studio/install/test_rocm_support.py | 65 +- tests/studio/install/test_selection_logic.py | 48 -- tests/studio/load_freeze/llama_server_shim.py | 4 +- .../load_freeze/test_load_orchestrator.py | 20 - tests/studio/playwright_chat_ime_i18n.py | 97 +-- tests/studio/run_real_mlx_smoke.py | 125 +-- tests/studio/test_auth_form_input_count.py | 22 +- tests/studio/test_cancel_id_wiring.py | 16 +- tests/studio/test_cli_studio_defaults.py | 6 +- .../test_composer_rtl_bidi_attribute.py | 7 +- tests/studio/test_frontend_dep_removal.py | 42 +- tests/studio/test_hardware_dispatch_matrix.py | 76 +- tests/studio/test_is_mlx_dispatch_gate.py | 18 +- .../test_stream_cancel_registration_timing.py | 77 +- tests/test_cli_export_unpacking.py | 47 +- tests/test_enforce_kwargs_spacing.py | 2 +- tests/test_finetune_last_n_layers.py | 13 +- tests/test_import_fixes_drift.py | 30 - tests/test_model_registry.py | 10 +- tests/test_multi_image_grpo_chunking.py | 6 +- tests/test_public_api_surface.py | 33 +- tests/test_raw_text.py | 4 +- tests/test_studio_install_workspace_guard.py | 23 +- tests/utils/aime_eval.py | 42 +- tests/utils/cleanup_utils.py | 34 +- tests/utils/data_utils.py | 36 +- tests/utils/generate_dataset_with_none.py | 35 +- tests/utils/ocr_eval.py | 37 +- tests/utils/os_utils.py | 4 +- tests/utils/perplexity_eval.py | 17 +- tests/utils/run_none_detect_tests.py | 21 +- tests/utils/test_packing.py | 2 +- tests/utils/test_q_galore.py | 9 +- tests/version_compat/_fetch.py | 33 +- .../test_bitsandbytes_pinned_symbols.py | 17 +- .../test_peft_pinned_symbols.py | 85 +-- ...st_sentence_transformers_pinned_symbols.py | 42 +- .../test_trl_grpo_pinned_symbols.py | 44 -- .../test_extended_module_imports.py | 66 +- tests/vllm_compat/test_unsloth_zoo_imports.py | 76 +- unsloth-cli.py | 26 +- unsloth/__init__.py | 28 +- unsloth/_gpu_init.py | 32 +- unsloth/chat_templates.py | 19 +- unsloth/dataprep/raw_text.py | 29 +- unsloth/dataprep/synthetic.py | 3 +- unsloth/device_type.py | 7 +- unsloth/import_fixes.py | 65 +- unsloth/kernels/__init__.py | 2 +- unsloth/kernels/cross_entropy_loss.py | 7 +- unsloth/kernels/fast_lora.py | 7 +- unsloth/kernels/fp8.py | 79 +- unsloth/kernels/moe/autotune_cache.py | 27 +- .../moe/benchmark/benchmark_fused_moe.py | 20 +- unsloth/kernels/moe/grouped_gemm/interface.py | 33 +- .../moe/grouped_gemm/kernels/autotuning.py | 5 +- .../moe/grouped_gemm/kernels/backward.py | 5 +- .../moe/grouped_gemm/kernels/tuning.py | 5 +- .../reference/layers/llama4_moe.py | 15 +- .../reference/layers/qwen3_moe.py | 23 +- unsloth/kernels/moe/tests/moe_utils.py | 26 +- .../kernels/moe/tests/test_grouped_gemm.py | 8 +- unsloth/kernels/utils.py | 16 +- unsloth/models/__init__.py | 2 +- unsloth/models/_utils.py | 18 +- unsloth/models/falcon_h1.py | 7 +- unsloth/models/gemma2.py | 7 +- unsloth/models/glm4_moe.py | 47 +- unsloth/models/granite.py | 5 +- unsloth/models/llama.py | 100 +-- unsloth/models/loader.py | 18 +- unsloth/models/loader_utils.py | 43 +- unsloth/models/mistral.py | 10 +- unsloth/models/qwen3.py | 44 +- unsloth/models/qwen3_moe.py | 15 +- unsloth/models/rl.py | 82 +- unsloth/models/sentence_transformer.py | 160 +--- unsloth/models/vision.py | 6 +- unsloth/optimizers/q_galore_adamw.py | 35 +- unsloth/optimizers/q_galore_projector.py | 24 +- unsloth/registry/__init__.py | 7 +- unsloth/registry/registry.py | 6 +- unsloth/save.py | 22 +- unsloth/tokenizer_utils.py | 24 +- unsloth/trainer.py | 38 +- unsloth/utils/attention_dispatch.py | 27 +- unsloth/utils/hf_hub.py | 14 +- unsloth_cli/_tool_policy.py | 19 +- unsloth_cli/commands/studio.py | 42 +- unsloth_cli/commands/train.py | 6 +- unsloth_cli/config.py | 3 +- unsloth_cli/options.py | 4 +- .../test_studio_run_short_alias_clashes.py | 6 +- 311 files changed, 4052 insertions(+), 8662 deletions(-) diff --git a/scripts/check_frontend_dep_removal.py b/scripts/check_frontend_dep_removal.py index 3ec4e9037f..b95c4ca7f6 100644 --- a/scripts/check_frontend_dep_removal.py +++ b/scripts/check_frontend_dep_removal.py @@ -43,7 +43,7 @@ DEP_FIELDS = ( "optionalDependencies", ) -# Sources where seeing a package name does NOT count as usage. +# Files where seeing a package name does NOT count as usage. EXPECTED_NOISE_FILES = { "studio/frontend/package.json", "studio/frontend/package-lock.json", @@ -51,17 +51,15 @@ EXPECTED_NOISE_FILES = { "studio/backend/core/data_recipe/oxc-validator/package-lock.json", } -# Only quoted-string occurrences in these file types can be module specifiers. +# File types where a quoted string can be a module specifier. JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$") -# Files where JS-syntactic import patterns (static/dynamic/require/re-export) -# could be a real module reference. Markdown gets a separate gate (.mdx is +# Files where JS import patterns could be a real module reference (.mdx is # real ESM; .md code fences are not). SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$") STYLE_EXT = re.compile(r"\.(css|scss|sass)$") HTML_EXT = re.compile(r"\.(html|htm)$") TS_LIKE_EXT = re.compile(r"\.(ts|tsx|mts|cts|mdx)$") -# Files where a removed package's CLI binary could be invoked (npx, bunx, -# yarn dlx, pnpm exec, or a bare `pkg --flag` shell call). +# Files where a removed package's CLI binary could be invoked. COMMAND_LIKE_EXT = re.compile(r"(\.(ya?ml|sh|ps1|bat)$|(^|/)Dockerfile[^/]*$)") GREP_INCLUDES = [ @@ -102,7 +100,7 @@ GREP_EXCLUDES = [ "--exclude-dir=venv", ] -# A pip-installed playwright reference is the PyPI package, not npm. +# A pip-installed playwright ref is the PyPI package, not npm. PIP_PLAYWRIGHT = re.compile( r"(pip\s+install\s+['\"]?playwright" r"|python\s+-m\s+playwright" @@ -153,9 +151,8 @@ def all_decl_names(pkg: dict) -> set[str]: def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None: - """Walk up the nested node_modules chain from `parent_path` to find - where `name` actually resolves. Mirrors Node module resolution. - """ + """Walk up the nested node_modules chain from `parent_path` to find where + `name` resolves, mirroring Node module resolution.""" parts = parent_path.split("/node_modules/") for i in range(len(parts), 0, -1): prefix = "/node_modules/".join(parts[:i]) @@ -168,11 +165,8 @@ def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None def _deps_of(meta: dict) -> dict: - """Deps npm actually installs. Optional peers are skipped: npm only - installs them when another package declares the same dep, so for the - purpose of "is this package still reachable" they cannot keep a - removed top-level dep alive on their own. - """ + """Deps npm actually installs. Optional peers are skipped: they can't keep + a removed top-level dep reachable on their own.""" out = {} for field in ("dependencies", "optionalDependencies"): out.update(meta.get(field) or {}) @@ -185,10 +179,8 @@ def _deps_of(meta: dict) -> dict: def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]: - """BFS the lockfile dep graph starting from `head_pkg`'s top-level - declared deps. Returns the set of lockfile install paths that survive. - Stale lockfile entries (orphaned by the new package.json) are excluded. - """ + """BFS the lockfile dep graph from `head_pkg`'s top-level deps. Returns the + surviving install paths, excluding stale (orphaned) lockfile entries.""" pkgs = lock.get("packages", {}) if not pkgs: return set() @@ -215,23 +207,17 @@ def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]: def classify(pkg: str, file: str, content: str) -> str | None: """Return why `content` references `pkg`, or None. - `content` may span multiple lines (for multi-line imports/exports); - each pattern uses re.DOTALL where it matters. The bare-spec - regexes use a word-boundary check on the package name so that - `foobar` does not match `foo`. - - File-type gating: JS-syntactic patterns only fire on .ts/.tsx/.js/.jsx/ - .mjs/.cjs/.mdx files, so an `import x from "pkg"` snippet inside a - Python test fixture or a Markdown code block is not mistaken for a - real npm usage. CSS patterns only fire on .css/.scss/.sass. HTML - patterns only fire on .html/.htm. + `content` may span multiple lines (multi-line imports/exports use re.DOTALL). + Bare-spec regexes word-boundary the package name so `foobar` doesn't match + `foo`. File-type gating restricts JS patterns to .ts/.tsx/.js/.jsx/.mjs/ + .cjs/.mdx, CSS to .css/.scss/.sass, HTML to .html/.htm, so a snippet inside + a Python fixture or Markdown code block isn't mistaken for real npm usage. """ if file in EXPECTED_NOISE_FILES: return None esc = re.escape(pkg) - # Subpath gate: after the package name, the next char must be either - # the closing quote, `/`, or end-of-string. Prevents foo matching foobar. + # Subpath gate: pkg must be followed by quote, `/`, or end-of-string. sub = r"(?:/[^'\"`]*)?" flags_dotall = re.DOTALL | re.MULTILINE @@ -241,30 +227,22 @@ def classify(pkg: str, file: str, content: str) -> str | None: is_html = bool(HTML_EXT.search(file)) is_ts = bool(TS_LIKE_EXT.search(file)) - # If the file is none of script / style / html / json (which is the - # quoted-string fallback surface) and is not an mdx file, no classify - # rule applies. This is what gates out Python fixtures, Markdown code - # blocks, shell snippets, etc. + # Gate out Python fixtures, Markdown code blocks, shell snippets, etc. is_json = file.endswith(".json") or file.endswith(".jsonc") if not (is_script or is_style or is_html or is_json): return None - # CSS @import is checked first so it does not collide with the - # side-effect-import regex below. + # CSS @import first so it doesn't collide with side-effect-import below. if is_style and re.search(rf"@import\s+['\"]{esc}{sub}['\"]", content): return "css_import" - # Static imports: handle multi-line `import { ... } from "pkg"` by - # allowing arbitrary content (newlines included) between `import` - # and `from`. The non-greedy match plus the required `from` keeps - # this scoped to a single statement. + # Static imports, including multi-line `import { ... } from "pkg"`. if is_script and re.search( rf"(? str | None: # require / require.resolve if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content): return "require" - # Re-exports: `export * from "pkg"`, `export { x } from "pkg"`, - # `export type { Foo } from "pkg"`. Multi-line supported. + # Re-exports: `export * from`, `export { x } from`, `export type { Foo } from`. if is_script and re.search( rf"\bexport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]", content, flags_dotall, ): return "re_export" - # HTML script / link. Match the package name as a complete path - # segment bounded by a quote / `#` / `?` or a subpath `/`, so - # `/node_modules/foo-extra/...` is NOT treated as usage of `foo`. + # HTML script / link. Match pkg as a complete path segment so + # `/node_modules/foo-extra/...` is not treated as usage of `foo`. html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])" if is_html and re.search(rf"]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content): return "html_script" @@ -295,8 +271,7 @@ def classify(pkg: str, file: str, content: str) -> str | None: # new URL("pkg/...", import.meta.url) if is_script and re.search(rf"\bnew\s+URL\(\s*['\"]{esc}{sub}['\"]", content): return "new_url" - # CSS url(...). Accept quoted ("pkg/x") AND unquoted (pkg/x) variants, - # bounded by a path-segment lookahead so `pkg-extra` does not match. + # CSS url(...), quoted and unquoted, bounded so `pkg-extra` doesn't match. if is_style and re.search( rf"\burl\(\s*['\"]?(?:[^)'\"\s]+/)?{esc}(?:/[^)'\"`]*)?['\"]?\s*\)", content, @@ -309,20 +284,18 @@ def classify(pkg: str, file: str, content: str) -> str | None: if is_script and re.search(rf"@import\(\s*['\"]{esc}{sub}['\"]\s*\)", content): return "jsdoc_import" # Bare quoted-string fallback (config plugin lists, vite aliases, - # tsconfig paths, biome config plugin arrays, shadcn registries). + # tsconfig paths, biome plugin arrays, shadcn registries). if not JS_LIKE_EXT.search(file): return None - # Boundary: pkg must be followed by `'`, `"`, or `/` to avoid - # matching `foo` inside `foobar`. + # pkg must be followed by `'`, `"`, or `/` so `foo` doesn't match `foobar`. if re.search(rf"['\"]{esc}(?:['\"]|/)", content): return "string_literal" return None def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]: - """Return a list of warnings if package-lock.json's dep map - disagrees with package.json (i.e., npm install was not re-run). - """ + """Warn if package-lock.json's dep map disagrees with package.json + (i.e. npm install was not re-run).""" warnings = [] if not head_lock: return warnings @@ -350,10 +323,8 @@ def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]: def types_orphan_warnings(head_pkg: dict) -> list[str]: - """Flag @types/ deps where is no longer declared anywhere - in package.json. Removing X without also dropping @types/X leaves - dangling type packages. - """ + """Flag @types/ deps where is no longer declared in package.json, + which leaves dangling type packages.""" decl = set() for f in DEP_FIELDS: decl.update((head_pkg.get(f) or {}).keys()) @@ -361,9 +332,7 @@ def types_orphan_warnings(head_pkg: dict) -> list[str]: for name in decl: if not name.startswith("@types/"): continue - # @types/foo provides types for `foo` - # @types/foo-bar provides types for `foo-bar` - # @types/scope__pkg provides types for `@scope/pkg` + # @types/scope__pkg provides types for @scope/pkg. target = name[len("@types/") :] if "__" in target: scope, sub = target.split("__", 1) @@ -386,8 +355,7 @@ _PKG_JSON_SKIP_KEYS = { "bundledDependencies", } -# Top-level fields whose contents are never package references. We walk -# everything else recursively. +# Top-level fields whose contents are never package references. _PKG_JSON_OPAQUE_KEYS = { "browserslist", # browser queries "keywords", # free-form strings @@ -429,20 +397,12 @@ _PKG_JSON_OPAQUE_KEYS = { def package_json_extra_refs(pkg: dict, target: str) -> list[str]: - """Walk every key/value in package.json EXCEPT the dep declaration - blocks, and return citations for string values or dict keys that - equal `target` (or `target/subpath`). + """Walk package.json (except dep declaration blocks) and return citations + for string values or dict keys equal to `target` (or `target/subpath`). - Catches the patterns the public dep-checker tools commonly miss: - - `overrides` / `resolutions` / `pnpm.overrides` keys - - `pnpm.patchedDependencies` keys - - `peerDependenciesMeta` keys - - `prettier`: "@my/prettier-config" - - `eslintConfig.extends`: ["..."] / "..." - - `stylelint.extends` / `stylelint.plugins` - - `babel.presets` / `babel.plugins` - - `jest.preset` / `jest.setupFiles` / `jest.transform` - - `commitlint.extends`, `renovate.extends`, `remarkConfig.plugins` + Catches refs that public dep-checkers commonly miss: overrides/resolutions/ + pnpm.overrides keys, pnpm.patchedDependencies, peerDependenciesMeta, + prettier, eslintConfig.extends, stylelint, babel, jest, commitlint, etc. """ target_sub = target + "/" cites: list[str] = [] @@ -453,14 +413,11 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]: def walk(obj: object, path: str) -> None: if isinstance(obj, dict): for k, v in obj.items(): - # Skip top-level dep declaration fields entirely. if path == "" and k in _PKG_JSON_SKIP_KEYS: continue - # Top-level fields whose contents are never package refs. if path == "" and k in _PKG_JSON_OPAQUE_KEYS: continue - # Inside `overrides` / `resolutions` / etc., the KEY itself - # is a package reference. + # Inside overrides/resolutions/etc., the KEY is a package ref. if matches(k): cites.append(f"{path}.{k}" if path else k) walk(v, f"{path}.{k}" if path else k) @@ -476,9 +433,8 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]: def build_bin_to_pkg(head_lock: dict) -> dict[str, str]: - """Map a binary name (e.g. 'vite', 'tsc', 'eslint') to the package - that provides it. Built from each lockfile entry's `bin` field. - """ + """Map a binary name (e.g. 'vite', 'eslint') to its providing package, + from each lockfile entry's `bin` field.""" out: dict[str, str] = {} if not head_lock: return out @@ -497,42 +453,29 @@ def build_bin_to_pkg(head_lock: dict) -> dict[str, str]: _SCRIPT_TOKENIZE = re.compile(r"\s*(?:&&|\|\||;|\|(?!\|))\s*") -# Wrappers that delegate to a real CLI in the same shell word list. -# After stripping env prefixes and (optionally) `npx`/`pnpm exec`/`yarn dlx`/ -# `bunx`, if the leading token is one of these we advance past the -# wrapper's own flags and any further env-prefix tokens, then re-check. -# `cross-env` is the common one; `dotenv-cli` / `dotenvx` use `--` as a -# separator. Wrappers that operate on named npm-scripts (concurrently, -# npm-run-all, run-s, run-p, wireit, turbo, nx) intentionally aren't -# here -- they reference script names, not bin names, so the real bin -# is in the *target* script's chunk which we already tokenize. +# Wrappers that delegate to a real CLI in the same shell word list; we skip +# past them and their flags to find the wrapped bin. Script-name wrappers +# (concurrently, npm-run-all, turbo, nx) are excluded: they reference script +# names, so the real bin lives in the target script's chunk we already tokenize. _SCRIPT_WRAPPERS = {"cross-env", "dotenv", "dotenvx", "env-cmd"} _ENV_PREFIX_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=") def _next_real_bin(words: list[str], idx: int) -> str | None: - """Walk `words` from `idx`, peeling env-prefix tokens, the leading - package-manager runner (`npx`, `pnpm exec`, etc.), and the known - wrapper bins. Return the next token that looks like the real CLI - binary, or None if the chunk has nothing to look up. - - Recursion depth is bounded by the chunk's word count, so the loop - cannot run away on a pathological wrapper chain. - """ + """Walk `words` from `idx`, peeling env-prefix tokens, the package-manager + runner (npx, pnpm exec, etc.), and known wrapper bins. Return the next + real CLI binary, or None. Bounded by the chunk's word count.""" seen_wrappers: set[str] = set() while idx < len(words): - # 1. env-prefix run: `FOO=bar BAZ="a b" cmd ...`. shlex has - # already collapsed quoted values into one word, so this - # tokenizer is safe for them. + # 1. env-prefix run `FOO=bar BAZ="a b" cmd ...` (shlex pre-collapsed). while idx < len(words) and _ENV_PREFIX_RE.match(words[idx]): idx += 1 if idx >= len(words): return None first = words[idx] - # 2. Package-manager runner: `npx args`, `pnpm exec `, - # `yarn dlx `, `bunx `. Strip and continue (so the - # wrapped command goes through the same unwrap loop). + # 2. Package-manager runner (npx/pnpm exec/yarn dlx/bunx): strip and + # continue so the wrapped command re-enters the unwrap loop. if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words): idx += 1 continue @@ -540,21 +483,17 @@ def _next_real_bin(words: list[str], idx: int) -> str | None: idx += 2 continue - # 3. Wrapper bin (cross-env, dotenv, etc.). Skip the wrapper's - # own flags and any subsequent env-prefix tokens, then re-loop. + # 3. Wrapper bin (cross-env, dotenv): skip its flags and env prefixes. bin_token = first.removeprefix("./node_modules/.bin/").removeprefix("node_modules/.bin/") if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers: seen_wrappers.add(bin_token) idx += 1 - # cross-env / env-cmd: no flags; just more env-prefix tokens. - # dotenv / dotenvx: skip `-e ` style flags and the - # optional `--` separator before the wrapped command. + # dotenv/dotenvx use `-e ` flags and an optional `--`. while idx < len(words): tok = words[idx] if tok.startswith("-") and tok != "--": idx += 1 - # `-e .env` style: also skip the flag's argument - # when it does not look like another flag. + # `-e .env`: also skip the flag's argument. if ( idx < len(words) and not words[idx].startswith("-") @@ -572,18 +511,12 @@ def _next_real_bin(words: list[str], idx: int) -> str | None: def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, list[str]]: - """Return `{package_name: ['scripts.X: cmd', ...]}` listing every - package referenced via its bin name in package.json scripts. + """Return `{package_name: ['scripts.X: cmd', ...]}` for every package + referenced via its bin name in package.json scripts. - Each script value is split on shell separators (`&&`, `||`, `;`, - `|`). Within each chunk, `_next_real_bin()` unwraps env prefixes, - package-manager runners (`npx` / `pnpm exec` / `yarn dlx` / `bunx`), - and wrapper bins like `cross-env` / `dotenv` so that - `cross-env CI=1 biome check` correctly credits `biome` to its - declaring package. - - Tokenization uses shlex.split so quoted env values - (`FOO="a b" biome`) survive unbroken. + Each script is split on shell separators; `_next_real_bin()` unwraps env + prefixes, package-manager runners, and wrapper bins so `cross-env CI=1 + biome check` credits `biome`. Uses shlex.split so quoted env values survive. """ import shlex @@ -599,7 +532,7 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li try: words = shlex.split(chunk, posix = True) except ValueError: - # Unbalanced quotes -- fall back to plain split. + # Unbalanced quotes: fall back to plain split. words = chunk.split() if not words: continue @@ -613,11 +546,8 @@ def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, li def tsconfig_compiler_types_refs() -> set[str]: - """Read studio/frontend/tsconfig*.json and return the set of - package names referenced in compilerOptions.types arrays. These are - implicitly loaded by tsc and count as a real use even though they - have no explicit import. - """ + """Return package names in tsconfig*.json compilerOptions.types arrays. + These are implicitly loaded by tsc and count as real uses.""" out: set[str] = set() base = REPO_ROOT / "studio/frontend" for name in ("tsconfig.json", "tsconfig.app.json", "tsconfig.node.json"): @@ -635,26 +565,16 @@ def tsconfig_compiler_types_refs() -> set[str]: for t in types: if not isinstance(t, str): continue - # `vite/client` resolves to `vite` package. + # `vite/client` resolves to the `vite` package. pkg = t.split("/", 1)[0] if not t.startswith("@") else "/".join(t.split("/", 2)[:2]) out.add(pkg) return out def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: - """For every declared dep, classify whether it appears used. Returns - a dict with these categories: - - used: has at least one detected usage in src/, - config files, scripts.bin, package.json - field refs, or tsconfig types - - unused: no detected usage anywhere - - type_pkg_kept: @types/X where X is still declared - - type_pkg_orphan: @types/X where X is no longer declared - (or X is removed) -- candidate for removal - - Each entry is the package name. The categorisation is opinionated; - `unused` is a CANDIDATE list, not a guarantee. The caller should - verify before deletion. + """For every declared dep, classify usage into a dict of package-name lists: + used, unused, type_pkg_kept (@types/X with X declared), type_pkg_orphan + (@types/X with X gone). `unused` is a CANDIDATE list; verify before deletion. """ decl = all_decl_names(head_pkg) bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} @@ -683,11 +603,8 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: # Real-source-usage check hits = find_usage(name) used = bool(hits) - # CLI usage in shell / workflow / Dockerfile surfaces. Skip for - # `@types/*` packages because they never expose a CLI binary and - # the unscoped-tail bin name candidate would scan workflow files - # for the bare runtime name (a removed `@types/foo` would look - # for invocations of `foo`). + # CLI usage in shell/workflow/Dockerfile. Skipped for @types/* (no CLI + # binary; the bare-name bin candidate would false-match the runtime). if not used and not name.startswith("@types/") and find_command_usage(name): used = True # Bin scripts @@ -707,28 +624,15 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]: def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: - """Reverse check: find bare-specifier imports in studio/frontend/src - that don't correspond to any declared package.json dep. Catches the - case where someone adds an import but forgets the dep declaration. - Returns (file, line, spec) tuples. - - Match shapes covered: - import "pkg" - import Foo from "pkg" - import { Foo } from "pkg" - import type { Foo } from "pkg" - const x = require("pkg") - const x = await import("pkg") + """Reverse check: find bare-specifier imports in studio/frontend/src with + no matching package.json dep (import added but dep declaration forgotten). + Covers import/require/dynamic-import shapes. Returns (file, line, spec). """ decl = set() for f in DEP_FIELDS: decl.update((head_pkg.get(f) or {}).keys()) - # Also: anything tsconfig path-aliases (just '@/...' here) is internal. - # The capture group is the specifier; the leading alternation accepts - # any of: `from "..."`, bare side-effect `import "..."`, - # `import("..."), or `require("...")`. We exclude relative paths and - # the `@/` alias prefix by requiring the first char of the specifier - # to be neither `.` nor `/`. + # Exclude relative paths and the `@/` alias by requiring the specifier's + # first char to be neither `.` nor `/`. Capture group is the specifier. pattern = ( r"(?:\bfrom\s+|" r"\bimport\s+(?:\(\s*)?|" @@ -754,7 +658,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: file, ln, content = m.group(1), int(m.group(2)), m.group(3) for spec_match in re.finditer(pattern, content): spec = spec_match.group(1) - # Resolve to package name (strip subpath) + # Resolve to package name (strip subpath). if spec.startswith("@"): parts = spec.split("/", 2) pkg_name = "/".join(parts[:2]) if len(parts) >= 2 else spec @@ -762,7 +666,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]: pkg_name = spec.split("/", 1)[0] if pkg_name in decl: continue - # Internal aliases like '@/foo' or starts with builtin names + # Internal aliases like '@/foo' or builtin names. if pkg_name == "@": continue if pkg_name in { @@ -807,11 +711,10 @@ def _read_file(path: str) -> list[str]: def find_usage(pkg: str) -> list[Hit]: - """Return real usages of `pkg`. Filters pip-playwright separately. + """Return real usages of `pkg` (pip-playwright filtered separately). - For each filename returned by grep, also feed a multi-line window - around the matching line into classify() so multi-line imports - (`import {\n a\n} from "pkg"`) get picked up. + For each grep hit, also feed a multi-line window into classify() so + multi-line imports get picked up. """ rows = grep_repo(re.escape(pkg)) hits = [] @@ -822,10 +725,8 @@ def find_usage(pkg: str) -> list[Hit]: # Try the single-line classify first. kind = classify(pkg, file, content) if not kind: - # Multi-line window: a generous 25 lines above + the line + - # 25 below so Prettier's one-import-per-line formatting for - # 12-20+ named imports still includes the `import` keyword - # in the same window as the `from "pkg"` clause. + # Multi-line window (25 lines each side) so Prettier's + # one-import-per-line formatting still pairs `import` with `from`. lines = _read_file(file) lo = max(0, lineno - 26) hi = min(len(lines), lineno + 25) @@ -841,28 +742,21 @@ def find_usage(pkg: str) -> list[Hit]: def _candidate_bin_names(pkg: str) -> set[str]: - """Names a removed package's CLI could be invoked under in shell - scripts and workflow files. Most npm CLIs use the package name - (`vite`, `eslint`, `playwright`); scoped CLI packages commonly - expose an unscoped binary name (`@biomejs/biome` -> `biome`). - """ + """Bin names a removed package's CLI could be invoked under. Most npm CLIs + use the package name; scoped ones expose an unscoped bin (@biomejs/biome -> + biome).""" return {pkg, pkg.rsplit("/", 1)[-1]} def find_command_usage(pkg: str) -> list[Hit]: - """Find package CLI invocations in shell / workflow / Dockerfile - surfaces: `npx pkg`, `bunx pkg`, `pnpm exec pkg`, `yarn dlx pkg`, - or a bare `pkg --flag`. Returns Hit("command_bin"). - - Detection is bounded to COMMAND_LIKE_EXT files so a JS string that - happens to contain `npx foo` inside a TS test fixture is not - mistaken for a real invocation. + """Find package CLI invocations in shell/workflow/Dockerfile surfaces (npx, + bunx, pnpm exec, yarn dlx, or bare `pkg --flag`). Bounded to + COMMAND_LIKE_EXT so `npx foo` in a TS fixture isn't mistaken for real use. """ bins = sorted(_candidate_bin_names(pkg), key = len, reverse = True) esc_bins = "|".join(re.escape(b) for b in bins) - # grep ERE pattern (POSIX classes for whitespace/word boundaries). - # Build without f-strings to avoid f-string-vs-{} confusion with the - # POSIX `[[:space:]]` literals and trailing `})}` boundary class. + # grep ERE pattern. Built without f-strings to avoid clashing with the + # POSIX `[[:space:]]` literals. grep_pat = ( r"(^|[[:space:]:;&|(\[])" r"(npx[[:space:]]+|pnpm[[:space:]]+exec[[:space:]]+" @@ -894,10 +788,8 @@ def find_command_usage(pkg: str) -> list[Hit]: def types_target_name(pkg: str) -> str | None: - """Strip `@types/` prefix and decode the npm scope-encoding so the - return value matches the runtime package name. `@types/foo` -> `foo`, - `@types/foo__bar` -> `@foo/bar`. Returns None for non-@types packages. - """ + """Strip `@types/` and decode scope-encoding to the runtime package name + (`@types/foo__bar` -> `@foo/bar`). None for non-@types packages.""" if not pkg.startswith("@types/"): return None target = pkg[len("@types/") :] @@ -908,11 +800,8 @@ def types_target_name(pkg: str) -> str | None: def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]: - """For a removed `@types/X`, find usages of `X` itself: explicit - `/// `, `tsconfig.compilerOptions.types: ["X"]`, - and runtime `import "X"` shapes. The whole point of `@types/X` is to - type one of those; if any are present, the type package must stay. - """ + """For a removed `@types/X`, find usages of `X` itself (triple-slash + reference, tsconfig types, runtime import). If any exist, @types/X stays.""" target = types_target_name(pkg) if target is None: return [] @@ -997,10 +886,9 @@ def main() -> int: return 2 head_lock = read_pkg_file(head_lock_path) - # Base lockfile is best-effort. We use it only to recover the - # bin -> package mapping for packages the PR is removing -- so a - # `scripts.biome:check` cite still fires when `@biomejs/biome` is - # being dropped and the head lockfile no longer has it. + # Base lockfile is best-effort: only used to recover the bin -> package + # mapping for packages the PR removes, so a scripts.biome cite still fires + # when @biomejs/biome is dropped from the head lockfile. if args.base_lock: base_lock_path = Path(args.base_lock) base_lock = read_pkg_file(base_lock_path) if base_lock_path.exists() else {} @@ -1011,9 +899,8 @@ def main() -> int: head_names = all_decl_names(head_pkg) removed = sorted(base_names - head_names) - # All hygiene checks compute up front so they can run on both the - # removal-present and removal-empty paths (so `--strict` actually - # fails when only hygiene issues exist). + # Hygiene checks compute up front so they run on both the removal-present + # and removal-empty paths (so --strict fails on hygiene-only issues). sync_warns = lockfile_root_sync(head_pkg, head_lock) types_warns = types_orphan_warnings(head_pkg) missing_imports = find_imports_without_decl(head_pkg) @@ -1074,12 +961,9 @@ def main() -> int: print() reachable_paths = reachable_from_head(head_pkg, head_lock) if head_lock else set() - # bin -> package map: start from the head lockfile, then layer the - # base lockfile's entries on top for packages this PR is removing. - # A correct removal updates the head lockfile to drop node_modules/foo, - # so build_bin_to_pkg(head_lock) loses the mapping; we recover it - # from the base lockfile so `scripts.biome:check` still flags as a - # usage when `@biomejs/biome` is being dropped. + # bin -> package map from the head lockfile, layering base-lockfile entries + # for removed packages so scripts.biome still flags when @biomejs/biome is + # dropped (head lockfile no longer maps it). bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {} base_bin_to_pkg = build_bin_to_pkg(base_lock) if base_lock else {} removed_set = set(removed) @@ -1091,9 +975,8 @@ def main() -> int: def reachable_install_paths(name: str) -> tuple[str | None, list[str]]: """Return (top_level_path, nested_paths). top_level is what bare - `import "name"` from src/ actually resolves to; nested copies are - only visible inside the parent package that nested them. - """ + `import "name"` resolves to; nested copies are only visible inside + their parent package.""" top = f"node_modules/{name}" top_path = top if top in reachable_paths else None nested = sorted( @@ -1106,8 +989,7 @@ def main() -> int: hits = find_usage(name) # CLI invocations in shell scripts / workflows / Dockerfiles. hits.extend(find_command_usage(name)) - # @types/X is "used" if X is referenced as a type or as a - # runtime import elsewhere in the repo. + # @types/X is "used" if X is referenced as a type or runtime import. hits.extend(find_types_runtime_usage(name, tsc_types)) for cite in script_refs.get(name, []): hits.append(Hit("studio/frontend/package.json", 0, "script_bin", cite)) @@ -1115,9 +997,8 @@ def main() -> int: hits.append(Hit("studio/frontend/package.json", 0, "pkg_json_field", cite)) top, nested = reachable_install_paths(name) importable_top_level = top is not None - # Source imports of bare specifier `name` resolve ONLY to top-level - # node_modules/. Nested copies under another package are - # invisible to src/ files. + # Bare specifier `name` resolves ONLY to top-level node_modules/; + # nested copies are invisible to src/ files. if hits and not importable_top_level: status = "FAIL" elif hits and importable_top_level: diff --git a/scripts/check_new_install_scripts.py b/scripts/check_new_install_scripts.py index 2beaada0bf..604d9b9f90 100644 --- a/scripts/check_new_install_scripts.py +++ b/scripts/check_new_install_scripts.py @@ -4,33 +4,18 @@ """Diff two `package-lock.json` files and flag NEW install-script deps. -A package with `"hasInstallScript": true` runs `preinstall` / `install` / -`postinstall` lifecycle hooks every time `npm ci` lays it down. Every -npm supply-chain compromise of the last 18 months (Shai-Hulud, -TanStack, axios-style, ArmorCode hijacks) leveraged exactly this lever: -the attacker publishes a new malicious version of a dep we already -trust, and the post-install hook runs the next time CI installs. +A `"hasInstallScript": true` package runs preinstall/install/postinstall +hooks on every `npm ci` -- the lever behind recent npm supply-chain +compromises (attacker publishes a malicious version of a trusted dep). +This refuses to land a newly-introduced install-script dep without a +maintainer eyeball; pre-existing ones are not re-flagged. -This scanner refuses to allow a newly-introduced install-script dep to -land without a maintainer eyeball on the lifecycle script body. -Existing install-script deps are NOT re-flagged -- if `node-gyp` has -been in the lockfile since day one, it's not part of this PR's threat -model. Only new entries are surfaced. +Supports lockfileVersion 1 (recursive `dependencies`) and 2/3 (flat +`packages` with `node_modules/.../node_modules/...` nesting). For each +new entry we best-effort fetch the registry metadata to recover the +postinstall command body; the finding is still emitted if unreachable. -Supports lockfileVersion 1 (`dependencies` key, recursive), 2 and 3 -(flat `packages` key with `node_modules//node_modules/` nesting -for transitive entries). For each NEW install-script package we -attempt a stdlib-only fetch of -`https://registry.npmjs.org//` to recover the actual -postinstall command body. If the network is blocked we still emit the -finding -- the lifecycle command body is informational, not -load-bearing. - -Exit codes -========== - 0 no newly-added install-script deps - 1 one or more newly-added install-script deps; listed on stderr - 2 internal error (missing lockfile, malformed JSON, etc.) +Exit codes: 0 = none; 1 = one or more (on stderr); 2 = internal error. """ from __future__ import annotations @@ -68,21 +53,14 @@ class Finding: ) -# ───────────────────────────────────────────────────────────────────── # Lockfile parsing. -# ───────────────────────────────────────────────────────────────────── def _strip_nm_prefix(key: str) -> str: - """Convert a v2/v3 `packages` key into a bare package name. - - `node_modules/foo` -> `foo`; `node_modules/foo/node_modules/bar` -> - `bar`. The empty key (`""`) is the project root and returns "". - """ + """Convert a v2/v3 `packages` key into a bare package name (leaf after last `node_modules/`).""" if not key: return "" - # Use the LAST `node_modules/` segment so transitives map to their - # leaf name, matching how npm install resolves a postinstall. + # LAST node_modules/ segment so transitives map to their leaf name. marker = "node_modules/" idx = key.rfind(marker) if idx == -1: @@ -91,14 +69,9 @@ def _strip_nm_prefix(key: str) -> str: def _collect_install_script_entries(lock: dict) -> dict[str, str]: - """Walk a parsed lockfile and return {package_name: version} for - every entry with `hasInstallScript: true` (v2/v3) OR a - non-empty `scripts.preinstall|install|postinstall` (v1). + """Return {name@version: name} for entries with hasInstallScript (v2/v3) or a lifecycle script (v1). - The same package may appear at multiple versions in a single - lockfile (de-duplicated copies under different parents); we key by - `name@version` so we don't lose either copy. Returns a dict keyed - by `name@version` -> the same string for convenience. + Keyed by name@version so dup copies at different versions aren't lost. """ seen: dict[str, str] = {} version = lock.get("lockfileVersion") @@ -118,10 +91,7 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]: ver = entry.get("version") or "" seen[f"{name}@{ver}"] = name - # v1 also embeds a `dependencies` tree; v2/v3 carry both for - # backwards-compat but `packages` is canonical for them. For v1 - # there is no `hasInstallScript` flag, so look for a non-empty - # `scripts.preinstall|install|postinstall` directly. + # v1 has no hasInstallScript flag; detect lifecycle scripts directly. def _walk_v1(deps: dict, depth: int = 0) -> None: if depth > 64 or not isinstance(deps, dict): return @@ -133,8 +103,6 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]: isinstance(scripts, dict) and scripts.get(hook) for hook in ("preinstall", "install", "postinstall") ) - # v1 also sets `requires` only on the parent, no flag, so - # the lifecycle-script presence is the only signal. if lifecycle: ver = entry.get("version") or "" seen[f"{name}@{ver}"] = name @@ -155,19 +123,11 @@ def _load_lockfile(path: Path) -> dict: raise ValueError(f"{path}: not valid JSON: {exc}") from exc -# ───────────────────────────────────────────────────────────────────── # Registry lookup for the postinstall command body (best-effort). -# ───────────────────────────────────────────────────────────────────── def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None: - """Return {hook: command} for any of preinstall / install / - postinstall published in the registry metadata for this name@ver. - - Returns None on any error (network blocked, 404, malformed JSON). - Never raises; the caller treats absence as "could not enrich, emit - finding anyway". - """ + """Return {hook: command} for lifecycle hooks in registry metadata; None on any error (never raises).""" safe_name = urllib.parse.quote(name, safe = "@/") url = f"{REGISTRY_BASE}{safe_name}/{urllib.parse.quote(version)}" try: @@ -190,9 +150,7 @@ def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None: return keep or None -# ───────────────────────────────────────────────────────────────────── # Diff. -# ───────────────────────────────────────────────────────────────────── def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: @@ -203,7 +161,6 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: if key in base: continue # pre-existing install-script dep; not in scope name = head[key] - # key is "name@version"; rsplit("@", 1) handles scoped names. version = key[len(name) + 1 :] if key.startswith(name + "@") else "" scripts = _fetch_registry_scripts(name, version) if scripts: @@ -226,9 +183,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]: return findings -# ───────────────────────────────────────────────────────────────────── # CLI. -# ───────────────────────────────────────────────────────────────────── def main(argv: list[str] | None = None) -> int: diff --git a/scripts/enforce_kwargs_spacing.py b/scripts/enforce_kwargs_spacing.py index fdef950d7f..dc4a6d6821 100755 --- a/scripts/enforce_kwargs_spacing.py +++ b/scripts/enforce_kwargs_spacing.py @@ -20,13 +20,8 @@ from pathlib import Path def _atomic_write_text(path: Path, data: str, encoding: str) -> None: - """Write ``data`` to ``path`` atomically. - - Stages a tmp file in the same directory (so it's on the same - filesystem as the destination), fsyncs, then `os.replace`s into - place. A crash mid-write therefore leaves either the previous - content or the fully new content -- never a truncated source file. - """ + """Write ``data`` to ``path`` atomically via same-dir tmp + fsync + os.replace, + so a crash mid-write leaves either the old or full new content, never a truncation.""" dirpath = str(path.parent) or "." fd, tmp_path = tempfile.mkstemp(prefix=".kwargs_fix.", dir=dirpath) try: @@ -142,7 +137,7 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]: lines[start] = segment if segment.strip() else "" continue - # Defensive fall-back for unexpected multi-line 'pass'. + # Fall-back for unexpected multi-line 'pass'. prefix = lines[start][: node.col_offset] lines[start] = prefix if prefix.strip() else "" for idx in range(start + 1, end): @@ -164,15 +159,12 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]: def remove_blank_after_short_import(text: str) -> tuple[str, bool]: - """Drop blank line(s) after an import block in a *small* nested suite. + """Drop blank line(s) after an import block in a small nested suite. - Inside an indented suite of <= 3 statements (function/try/if/with/etc., never - module level), when a run of consecutive ``import`` / ``from ... import`` - statements is directly followed -- across one or more blank lines and nothing - else -- by another statement in the same suite, remove those blank lines so - the import sits next to the code that uses it. A comment in the gap blocks the - rule. Removing blank lines never changes the AST, so this is always - semantics-preserving. + In an indented suite of <= 3 statements (never module level), when consecutive + imports are followed across blank lines (nothing else) by another statement, + remove those blanks. A comment in the gap blocks the rule. Removing blank lines + never changes the AST. """ try: tree = ast.parse(text) @@ -226,12 +218,11 @@ _DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]: - """Map the line of each def keyword to (param count, has-any-default). + """Map each def keyword line to (param count, has-any-default). - One def per line, so the line is a stable key. ``*`` / ``/`` markers are not - parameters and are not counted. A default exists if any positional default - is present or any keyword-only default is not ``None`` (a ``None`` entry in - ``kw_defaults`` means a required keyword-only arg). + ``*`` / ``/`` markers aren't counted. A default exists if any positional default + is present or any keyword-only default is not ``None`` (``None`` in ``kw_defaults`` + means a required keyword-only arg). """ out: dict[int, tuple[int, bool]] = {} for node in ast.walk(tree): @@ -250,20 +241,12 @@ def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]: def normalize_def_trailing_comma(text: str) -> tuple[str, bool]: - """Force a def / async-def signature one-per-line iff it has >= 3 parameters - AND at least one default value; otherwise keep it collapsible. + """Force a def signature one-per-line iff >= 3 params AND a default; else collapsible. - Rationale: signatures with defaults read better one parameter per line, but - only once they are non-trivial (< 3 params always stay on one line). A - signature with >= 3 params and a default gets a magic trailing comma added - (ruff then wraps it one-per-line regardless of length); every other - signature has its trailing comma stripped so ruff collapses it onto one line - when it fits (and wraps a genuinely long one by length alone). - - Function-definition parameter lists only, never call sites or collection - literals. Parameter counts and defaults come from the AST. Run BEFORE ruff - format. Adding or removing a def trailing comma never changes the AST, which - is re-checked before returning. + A qualifying signature gets a magic trailing comma added (ruff wraps it + one-per-line); every other signature has its trailing comma stripped so ruff + collapses it when it fits. Def parameter lists only, never call sites or + collection literals. Run BEFORE ruff format. Never changes the AST (re-checked). """ try: tree = ast.parse(text) @@ -332,12 +315,10 @@ def normalize_def_trailing_comma(text: str) -> tuple[str, bool]: def _split_string_token(s: str) -> tuple[str, str, str] | None: - """Split a string literal's source into (prefix, quote, body). + """Split a string literal source into (prefix, quote, body). - ``prefix`` is the letters before the opening quote (``r``/``f``/``b``/``u`` - in any case/order), ``quote`` is the opening delimiter (``'``, ``"``, - ``'''`` or ``\"\"\"``) and ``body`` is everything between the delimiters. - Returns ``None`` if ``s`` is not a recognizable string literal. + ``prefix`` is the letters before the opening quote, ``quote`` the delimiter, + ``body`` everything between. ``None`` if not a recognizable string literal. """ i = 0 while i < len(s) and s[i] not in ("'", '"'): @@ -393,15 +374,12 @@ def _string_pieces( def _merge_string_run(pieces: list[tuple[str, str]]) -> str | None: """Merge a run of adjacent string pieces into one literal's source text. - ``pieces`` is a list of ``(kind, raw_source)`` where kind is ``"str"`` or - ``"f"``. Rules: bytes are left side-by-side (return ``None``); a run with no - f-string merges plain/raw/unicode pieces sharing one prefix+quote by simple - body concatenation; a run mixing an f-string with at least one plain string - (and no bytes, no raw) folds into a single f-string -- f pieces keep their - bodies verbatim and plain pieces have their braces escaped (``{`` -> ``{{``). - Runs of only f-strings are left side-by-side. The caller re-checks the file - AST and drops the change if it differs, so any subtle case (e.g. ``\\N{...}``) - that this would mis-handle is caught and skipped. + ``pieces`` is ``(kind, raw_source)`` with kind ``"str"`` or ``"f"``. Bytes are + left side-by-side (``None``); a run with no f-string merges plain/raw/unicode + sharing one prefix+quote by body concatenation; a run mixing an f-string with a + plain string (no bytes, no raw) folds into one f-string with plain braces escaped. + Runs of only f-strings are left alone. Caller re-checks the AST and drops a + differing change, so subtle cases are caught. """ parsed = [] for kind, raw in pieces: @@ -455,13 +433,10 @@ def _fold_collapses( ) -> bool: """Whether an f-string fold at ``row[c0:c1]`` -> ``merged`` is safe to apply. - Only ``assert`` statements wrap awkwardly when a message is folded: ruff - parenthesizes the *condition* once ``assert cond, msg`` no longer fits on one - line. For every other construct (call argument, ``raise``, assignment, ...) a - folded long message wraps acceptably, so the fold is always allowed. For an - ``assert`` the fold is allowed only when the statement is already one physical - line, or its estimated one-line length after folding fits the line length; - otherwise the message is left side-by-side. + Only ``assert`` wraps awkwardly when a message folds (ruff parenthesizes the + condition once it no longer fits one line); every other construct wraps + acceptably so is always allowed. An ``assert`` fold is allowed only if already + one line, or its estimated folded one-line length fits the line length. """ stmt = _enclosing_stmt(tree, row) if not isinstance(stmt, ast.Assert): @@ -483,15 +458,12 @@ def _fold_collapses( def merge_adjacent_string_literals(text: str) -> tuple[str, bool]: - """Merge a run of adjacent string literals on ONE physical line into a single - literal (the ``"a" "b"`` form ruff emits when it collapses an implicit - concatenation). Plain/raw/unicode runs merge by concatenation; a run mixing - an f-string with a plain string folds into one f-string (plain parts' braces - escaped) -- but only when the statement still fits on one line, so a long - message is left side-by-side rather than forcing the statement to re-wrap. - Runs of only f-strings, and bytes, are left side-by-side. The whole file's AST - is re-checked and the change is dropped if it would differ, so the transform - can never change meaning. + """Merge adjacent string literals on ONE physical line into a single literal. + + Plain/raw/unicode runs merge by concatenation; an f-string + plain string folds + into one f-string (plain braces escaped) only while the statement still fits one + line. Runs of only f-strings, and bytes, are left side-by-side. The file AST is + re-checked and a differing change dropped, so meaning never changes. """ try: toks = list(tokenize.generate_tokens(io.StringIO(text).readline)) @@ -552,17 +524,11 @@ def merge_adjacent_string_literals(text: str) -> tuple[str, bool]: def collapse_short_asserts(text: str) -> tuple[str, bool]: """Collapse a multi-line ``assert`` onto one line when it would fit. - An ``assert`` is often kept multi-line only by a magic trailing comma inside - a collection / call / tuple-message (``assert x == {\"a\": 1,}`` written - across lines). When the whole statement's estimated one-line length fits the - line length, strip those trailing commas (the comma before a ``)`` / ``]`` / - ``}``) so ruff joins it back onto one line on the following format pass. - - Run BEFORE ruff format. Skips any assert that contains a comment (a comment - forces ruff to keep it multi-line, which would oscillate). Stripping a - trailing comma is non-semantic except for a one-element tuple ``(x,)``; the - file AST is re-checked and any assert whose strip would change it is left - alone. + When the statement's estimated one-line length fits, strip the magic trailing + commas (comma before a closer) holding it open so ruff rejoins it. Run BEFORE + ruff format. Skips asserts with a comment (would oscillate). Stripping is + non-semantic except for a one-element tuple; AST is re-checked and changing + asserts left alone. """ try: tree = ast.parse(text) diff --git a/scripts/lint_workflow_triggers.py b/scripts/lint_workflow_triggers.py index 34868b7d68..0688f6c65c 100644 --- a/scripts/lint_workflow_triggers.py +++ b/scripts/lint_workflow_triggers.py @@ -4,37 +4,19 @@ """Refuse dangerous GitHub Actions trigger patterns at PR time. -Two patterns are banned outright, both of which powered the TanStack -GHSA-g7cv-rxg3-hmpx supply-chain compromise: +Bans patterns behind the TanStack GHSA-g7cv-rxg3-hmpx compromise: -1. `pull_request_target` -- runs a fork's workflow YAML against the - BASE repository's secrets and permissions. The fork can inject - arbitrary code into the base context. The TanStack worm used this - to land base-context execution from a fork PR. There is essentially - no safe use of this trigger for a public open-source project; - `pull_request` is the safe alternative. +1. `pull_request_target` -- runs a fork's workflow against the base + repo's secrets/permissions; use `pull_request` instead. +2. `workflow_run` chained to a PR-triggered workflow -- same trust + boundary problem one hop later (poisoned artifacts/caches run with + elevated permissions). +3. Cache keys shared between PR-triggered and publish/release/push + workflows -- a fork PR could poison a cache the publish workflow + restores. Partition the key namespaces. -2. `workflow_run` chained to a PR-triggered workflow -- carries the - same trust boundary problem one hop later. If a PR-triggered - workflow can poison artifacts/caches and a `workflow_run` trigger - fires off the result with elevated permissions, the attacker still - reaches the trusted context. - -3. Shared cache keys between PR-triggered workflows and publish / - release / push-triggered workflows. The TanStack worm poisoned the - Actions cache from a fork PR and the legitimate release workflow - then restored the poisoned cache. Cache keys must be partitioned - so that nothing a PR can write is ever read by a workflow that - holds secrets. - -Exit codes -========== - - 0 no findings - 1 one or more findings; stderr lists each with file path - -Run from repo root: - python3 scripts/lint_workflow_triggers.py +Exit codes: 0 = no findings, 1 = findings (listed on stderr). +Run from repo root: python3 scripts/lint_workflow_triggers.py """ from __future__ import annotations diff --git a/scripts/lockfile_supply_chain_audit.py b/scripts/lockfile_supply_chain_audit.py index 5b503dcffe..66b48c094d 100644 --- a/scripts/lockfile_supply_chain_audit.py +++ b/scripts/lockfile_supply_chain_audit.py @@ -5,64 +5,21 @@ """Lockfile supply-chain audit for the Studio frontend and Tauri shell. Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a -lockfile contains patterns that indicate the kind of supply-chain -injection seen in the npm Shai-Hulud waves and the cargo -crates.io brand-squat attempts. +lockfile contains patterns indicating supply-chain injection (npm +Shai-Hulud waves, cargo crates.io brand-squats). -What it checks -============== +Checks package-lock.json (lockfileVersion 2/3): `resolved` URL must be +the npm registry (direct git/github/file refs are the injection vector); +`integrity` SHA must be present; known IOC substrings grepped from the +body. Checks Cargo.lock: `source` must be the crates.io registry index; +known cargo IOC substrings. -studio/frontend/package-lock.json (lockfileVersion 2 or 3): +Exit codes: 0 = clean (or skip env var set to a justification >=5 chars, +not '1'/'true'); 1 = findings; 2 = internal error. - 1. `resolved` URL origin. Every entry must resolve through - `https://registry.npmjs.org/`. Direct GitHub-hosted dependencies - (`git+ssh://`, `git+https://`, `github:owner/repo#sha`, - `file:`, `http://`) are refused -- npm's TanStack incident used - exactly this vector to land an unaudited GitHub commit hash as - an optional dependency. - - 2. `integrity` field presence. Every non-workspace entry must carry - an `integrity` SHA. A missing integrity means the registry can - swap the tarball after lockfile generation and CI will not - notice. - - 3. Known IOC strings. A hardcoded set of indicator-of-compromise - substrings is grepped across the entire lockfile body (file - names, dependency keys, URLs). The list is updated as new - campaigns surface. Catching one means the local install was - about to pull a publicly-known malicious release. - -studio/src-tauri/Cargo.lock: - - 4. `source` field origin. Every entry with a `source` must point at - `registry+https://github.com/rust-lang/crates.io-index`. Direct - git sources (`git+https://...`) and `path+...` for cross-crate - paths warrant manual review and are flagged. - - 5. Known cargo IOC strings. Same idea as (3), separate list. - -Exit codes -========== - - 0 no findings, or an opt-out env var (UNSLOTH_LOCKFILE_AUDIT_SKIP) - is set to a justification string (>=5 chars, not '1'/'true'/etc). - A value like '1' or 'true' is now REJECTED loudly and the audit - runs normally - 1 one or more findings; stderr lists them with file path and line - number where derivable - 2 internal error (missing dependency, malformed JSON, etc.) - -Operational stance -================== - -This scanner only PARSES the lockfiles -- it never executes anything -in them, never resolves anything against the network. Safe to run -ahead of every `npm ci`. The IOC list is short by design; this -complements (not replaces) `npm audit`, OSV-Scanner, and the -advisory-DB pipeline in `.github/workflows/security-audit.yml`. The -shape of the catch is "we refuse to proceed because the lockfile -itself is shaped wrong", which fires before any third-party install -script gets a chance to run on the runner. +Only PARSES the lockfiles, never executes or networks. Complements (not +replaces) `npm audit` / OSV-Scanner / the advisory-DB pipeline. Fires +before any third-party install script runs on the runner. """ from __future__ import annotations @@ -77,14 +34,9 @@ from pathlib import Path REPO_ROOT = Path(__file__).resolve().parents[1] -# ───────────────────────────────────────────────────────────────────── -# Known IOC strings (case-sensitive substring match). -# ───────────────────────────────────────────────────────────────────── -# -# Keep these short and FACTUAL. Each entry is tied to a public advisory -# and is the literal string an attacker would have to embed for the -# attack to work. Adding speculative or generic patterns here would -# generate false positives on dependency upgrades. +# Known IOC strings (case-sensitive substring match). Each is tied to a +# public advisory; speculative/generic patterns would false-positive on +# upgrades. NPM_IOC_STRINGS: tuple[str, ...] = ( # Shai-Hulud TanStack wave -- May 11, 2026 (GHSA-g7cv-rxg3-hmpx). "router_init.js", @@ -328,36 +280,22 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = { } CARGO_IOC_STRINGS: tuple[str, ...] = ( - # Reserved for future cargo-side incidents. Empty by default -- - # `source` origin check below catches the structural pattern. + # Empty by default; the `source` origin check catches the structural + # pattern. Reserved for future cargo-side incidents. ) -# ───────────────────────────────────────────────────────────────────── # Allowed lockfile origins. -# ───────────────────────────────────────────────────────────────────── NPM_REGISTRY_PREFIX = "https://registry.npmjs.org/" - -# Tarballs are also fetched from this mirror on some GH Actions cached -# runs (npm rewrites the resolved URL on cache hit). Allow either. NPM_REGISTRY_PREFIXES_ALLOWED: tuple[str, ...] = (NPM_REGISTRY_PREFIX,) CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index" -# ───────────────────────────────────────────────────────────────────── -# Cargo non-registry source allowlist. -# ───────────────────────────────────────────────────────────────────── -# -# Each entry is `(crate_name, exact_source_string)`. The crate must -# match by name AND the source must match the full pinned-SHA string -# verbatim. Bumping the commit SHA forces a re-review here: the -# scanner fires until the new SHA is appended. -# -# Studio's Tauri shell pulls `fix-path-env` directly from -# tauri-apps/fix-path-env-rs because the crate is not published to -# crates.io. The pinned commit (c4c45d5) was reviewed at the time it -# landed; future bumps need explicit approval. +# Cargo non-registry source allowlist: `(crate_name, exact_source_string)`. +# Both must match verbatim; bumping the pinned SHA forces a re-review. +# Studio's Tauri shell pulls `fix-path-env` from git because it is not +# published to crates.io; commit c4c45d5 was reviewed when it landed. CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( ( "fix-path-env", @@ -367,11 +305,6 @@ CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = ( ) -# ───────────────────────────────────────────────────────────────────── -# Finding container. -# ───────────────────────────────────────────────────────────────────── - - class Finding: __slots__ = ("path", "package", "kind", "detail") @@ -390,27 +323,19 @@ class Finding: def _gha_escape(text: str) -> str: - """Escape a string for use in a GitHub Actions `::warning::` / - `::error::` workflow command message. GH Actions truncates - annotation messages at the first newline unless `\\n` is - escaped as `%0A`; carriage returns and the percent sign need - matching escapes per the workflow-commands spec. Order matters: - `%` must be replaced first so the subsequent `%0A` / `%0D` - sequences are not double-encoded. + """Escape a string for a GH Actions `::warning::`/`::error::` message. + + GH Actions truncates at the first newline unless `\\n`/`\\r` are + escaped as `%0A`/`%0D`. `%` must be replaced first to avoid + double-encoding the subsequent escapes. """ return text.replace("%", "%25").replace("\r", "%0D").replace("\n", "%0A") -# ───────────────────────────────────────────────────────────────────── -# package-lock.json audit. -# ───────────────────────────────────────────────────────────────────── - - def audit_npm_lockfile(path: Path) -> list[Finding]: findings: list[Finding] = [] if not path.exists(): - # A missing requested lockfile is a config error, not a clean - # audit; surface it so a deleted default cannot pass silently. + # Missing lockfile is a config error, not a clean audit. findings.append( Finding( path = str(path), @@ -427,8 +352,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: try: raw = path.read_text(encoding = "utf-8") except OSError as exc: - # Permission denied, is-a-directory, broken-pipe etc. -- surface - # as a finding instead of crashing CI with a raw traceback. + # Surface as a finding instead of crashing CI with a traceback. findings.append( Finding( path = str(path), @@ -464,9 +388,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: packages = lock.get("packages") or {} for key, entry in packages.items(): - # The empty key "" is the project root; workspace entries use - # keys like "node_modules/foo" or "studio/frontend/sub-pkg". - # Skip the project root (it has no `resolved`). + # Empty key "" is the project root (no `resolved`); skip it. if key == "": continue if entry.get("link"): @@ -474,12 +396,8 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: continue resolved = entry.get("resolved") - # Entries living inside another package's `node_modules/` - # tree are bundled fold-ins -- the parent's tarball ships - # their source verbatim and the parent's `integrity` covers - # the whole subtree. npm represents them in lockfileVersion 3 - # as nested entries with no `resolved` and no `integrity` of - # their own. Treat them as transparent to this audit. + # Entries nested in another package's node_modules are bundled + # fold-ins covered by the parent's integrity; treat as transparent. nested = key.count("/node_modules/") >= 1 # 1. resolved-URL origin. @@ -545,12 +463,10 @@ def audit_npm_lockfile(path: Path) -> list[Finding]: ) ) - # 4. Known IOC strings: scan the raw file body so we hit fields the - # structural pass above doesn't enumerate (scripts, optional - # dependencies, etc.). Cheap and complete. + # 4. Known IOC strings: scan the raw body to catch fields the + # structural pass doesn't enumerate (scripts, optional deps, etc.). for ioc in NPM_IOC_STRINGS: if ioc in raw: - # Best-effort line number lookup. line_no = _first_line_containing(raw, ioc) findings.append( Finding( @@ -575,14 +491,7 @@ def _first_line_containing(text: str, needle: str) -> int | None: return None -# ───────────────────────────────────────────────────────────────────── -# Cargo.lock audit. -# ───────────────────────────────────────────────────────────────────── - - -# Cargo.lock is TOML; parse with stdlib tomllib (Python 3.11+). The -# studio's Tauri shell already requires a modern toolchain so this is -# always available where CI runs. +# Cargo.lock is TOML; parsed with stdlib tomllib (Python 3.11+). _PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$") @@ -706,24 +615,15 @@ def audit_cargo_lockfile(path: Path) -> list[Finding]: return findings -# ───────────────────────────────────────────────────────────────────── -# CLI. -# ───────────────────────────────────────────────────────────────────── - - # Finding kinds split into BLOCKING vs ADVISORY for the default run mode. -# Blocking findings come from public supply-chain attack indicators (a -# version we know is malicious, a string an attacker would have to embed -# for an attack to work). Advisory findings are structural lockfile -# anomalies (missing integrity, non-default registry, etc.) -- they -# WARN the maintainer but do not block merges. Pass --strict to make -# every finding blocking (PR-5479-style behavior for opt-in adopters). +# Blocking = public attack indicators (known-malicious version, IOC +# string). Advisory = structural anomalies that warn but don't block. +# --strict makes every finding blocking. BLOCKING_KINDS: frozenset[str] = frozenset( { "blocked-known-malicious", "known-ioc-string", - # Internal-failure kinds: a structurally broken lockfile MIGHT - # be hiding a real attack, so we keep these blocking too. + # A structurally broken lockfile might hide a real attack. "malformed-lockfile", "missing-lockfile", "unreadable-lockfile", @@ -779,15 +679,9 @@ def main(argv: list[str] | None = None) -> int: ) args = parser.parse_args(argv) - # SF4: require a real justification (e.g. JIRA ticket id) for the - # skip env var. Treat the trivially-set values ("1", "true", "yes", - # "on", empty) as INVALID -- they look like accidental flips and - # silently bypassed the supply-chain audit. A valid value is a - # non-empty string >=5 chars after stripping that does not match - # any of the boolean-shaped tokens above. An invalid value emits a - # loud GitHub Actions warning to stderr and FALLS THROUGH to run - # the audit normally (fail-safe). A valid value emits a warning - # naming the reason and skips with rc=0 (compat). + # Require a real justification (>=5 chars, not a boolean-shaped token) + # for the skip env var. An invalid value warns and falls through to + # run the audit (fail-safe); a valid one warns and skips with rc=0. _skip_raw = os.environ.get("UNSLOTH_LOCKFILE_AUDIT_SKIP") if _skip_raw is not None: _skip = _skip_raw.strip() @@ -809,8 +703,7 @@ def main(argv: list[str] | None = None) -> int: return 0 root = Path(args.root).resolve() - # Explicit --npm-lockfile/--cargo-lockfile scopes the scan to those - # paths; defaults apply only to the no-args CI invocation. + # Explicit flags scope the scan; defaults apply only to no-args CI. _user_explicit = args.npm_lockfile is not None or args.cargo_lockfile is not None if _user_explicit: npm_paths = [root / p for p in (args.npm_lockfile or ())] @@ -835,11 +728,9 @@ def main(argv: list[str] | None = None) -> int: ) return 0 - # Split findings into blocking (known-malicious / IOC / structurally - # broken) and advisory (everything else, e.g. missing integrity on a - # registry-published tarball). In default mode advisory findings are - # printed but do not change the exit code; --strict treats every - # finding as blocking. + # Split into blocking (known-malicious / IOC / structurally broken) + # and advisory (everything else). Default mode prints advisories + # without changing the exit code; --strict makes all blocking. blocking = [f for f in all_findings if f.kind in BLOCKING_KINDS] advisory = [f for f in all_findings if f.kind not in BLOCKING_KINDS] @@ -854,12 +745,8 @@ def main(argv: list[str] | None = None) -> int: file = sys.stderr, ) for f in advisory: - # Surface in GitHub Actions UI as a warning annotation when run - # under Actions; harmless prefix elsewhere. GH Actions - # truncates annotation messages at the first newline unless - # newlines are escaped as `%0A`, so the full multi-line - # Finding (kind + path + package + detail) only renders in - # the UI after _gha_escape collapses it onto one line. + # GH Actions warning annotation; _gha_escape collapses the + # multi-line Finding onto one line so it renders fully in the UI. print(f"::warning::{_gha_escape(str(f))}", file = sys.stderr) print(file = sys.stderr) @@ -876,9 +763,7 @@ def main(argv: list[str] | None = None) -> int: file = sys.stderr, ) for f in blocking: - # Same %-encoding rationale as the advisory branch above: the - # GH Actions annotation is truncated at the first newline - # unless the message is escaped. + # Same %-encoding rationale as the advisory branch above. print(f"::error::{_gha_escape(str(f))}", file = sys.stderr) print(file = sys.stderr) print( diff --git a/scripts/notebook_to_python.py b/scripts/notebook_to_python.py index cc9dc0acda..4b64123d6a 100644 --- a/scripts/notebook_to_python.py +++ b/scripts/notebook_to_python.py @@ -22,19 +22,14 @@ import urllib.parse from pathlib import Path -# Hosts we are willing to fetch raw notebook JSON from. Anything else -# is rejected before `urlopen` so a typoed / hostile URL cannot pull -# code from arbitrary infrastructure. +# Allowlist of hosts for raw notebook fetches; anything else rejected before urlopen. _ALLOWED_NOTEBOOK_HOSTS = { "raw.githubusercontent.com", "gist.githubusercontent.com", } -# Shell metacharacters that imply the cell's `!cmd` line cannot be -# parsed as a flat argv. If any of these appears, `shlex.split` would -# either fail or, worse, silently strip the operator -- so we keep -# `shell=True` for that command and emit a review marker. +# Metacharacters that mean a `!cmd` line can't be a flat argv -> keep shell=True + review marker. _SHELL_METACHARS_RE = re.compile(r"\$\(|`|\|\||\||&&|>>?|< bool: def github_blob_to_raw(url: str) -> str: """Convert GitHub blob URL to raw URL.""" - # https://github.com/user/repo/blob/branch/path - # -> https://raw.githubusercontent.com/user/repo/branch/path - # Compare the parsed host exactly (not as a substring) so a URL - # like https://attacker.example.com/github.com/blob/... does NOT - # get rewritten to a github raw URL. Closes CodeQL alert - # py/incomplete-url-substring-sanitization. + # github.com/user/repo/blob/branch/path -> raw.githubusercontent.com/user/repo/branch/path + # Exact host match (not substring) so attacker.example.com/github.com/blob/... is not rewritten. parsed = urllib.parse.urlparse(url) if parsed.netloc != "github.com" or "/blob/" not in parsed.path: return url @@ -63,18 +54,12 @@ def github_blob_to_raw(url: str) -> str: def download_notebook(url: str) -> tuple[str, str]: """Download notebook from URL. Returns (content, filename).""" - # Convert blob URL to raw if needed raw_url = github_blob_to_raw(url) - # Extract filename from URL parsed = urllib.parse.urlparse(raw_url) filename = os.path.basename(urllib.parse.unquote(parsed.path)) - # Host allowlist. Refuse to fetch from anywhere the campaign IOC - # tables flag (or just anywhere we don't recognise). The blob->raw - # conversion above only emits `raw.githubusercontent.com`, so a - # rejection here means the caller hand-typed a URL pointing - # somewhere we don't trust. + # Host allowlist: refuse to fetch from anything we don't recognise. host = parsed.hostname if host not in _ALLOWED_NOTEBOOK_HOSTS: raise ValueError( @@ -82,7 +67,6 @@ def download_notebook(url: str) -> tuple[str, str]: f"{sorted(_ALLOWED_NOTEBOOK_HOSTS)}" ) - # Download print(f"Downloading {url}...") with urllib.request.urlopen(raw_url, timeout = 60) as response: content = response.read().decode("utf-8") @@ -97,29 +81,18 @@ def is_url(path: str) -> bool: def replace_colab_paths(source: str) -> str: """Replace Colab-specific /content/ paths with current working directory.""" - # Replace /content/ with f-string using _WORKING_DIR source = source.replace('"/content/', 'f"{_WORKING_DIR}/') source = source.replace("'/content/", "f'{_WORKING_DIR}/") return source def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> list[str]: - """Render a `!cmd` notebook line as one or more Python statements. + """Render a `!cmd` notebook line as Python statements. - When the command body is f-string-interpolated, contains shell - metacharacters, or spans multiple lines, falling back to - `shell=True` is the only correct option -- `shlex.split` would - either drop operators or fail outright. We surface that with a - `# WARNING: shell=True; reviewed for hostile input` comment so a - reviewer cannot miss it. - - Otherwise we emit `subprocess.run(shlex.split(cmd), shell=False)` - so the converted script is not a re-injection vector if the - notebook ever interpolates user-controlled data. - - `allow_shell` defaults to True at the CLI for backwards - compatibility. Setting it to False makes `shell=True` emission a - hard error (no surprise behaviour). + f-string interpolation, shell metacharacters, or multiline force + shell=True (shlex.split would drop operators), flagged with a + WARNING comment. Otherwise emit shell=False argv form. allow_shell + False makes shell=True emission a hard error. """ needs_f = needs_fstring(full_cmd) has_meta = bool(_SHELL_METACHARS_RE.search(full_cmd)) @@ -144,7 +117,6 @@ def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> lis stmt = f"{indent}subprocess.run({f_prefix}{full_cmd!r}, shell=True)" return [warn, stmt] - # Shell-safe argv form. return [f"{indent}subprocess.run(shlex.split({full_cmd!r}), shell=False)"] @@ -159,12 +131,10 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str: stripped = line.strip() indent = line[: len(line) - len(line.lstrip())] - # Skip %%capture if stripped.startswith("%%capture"): i += 1 continue - # Handle %%file magic if stripped.startswith("%%file "): filename = stripped[7:].strip() file_lines = [] @@ -178,7 +148,6 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str: result.append(f'{indent} _f.write("""{file_content}""")') continue - # Handle ! shell commands if stripped.startswith("!"): cmd_lines = [stripped[1:]] while cmd_lines[-1].rstrip().endswith("\\") and i + 1 < len(lines): @@ -308,21 +277,16 @@ def convert_notebook_to_script( content = f.read() source_name = source - # Generate output filename output_filename = filename.replace(".ipynb", ".py") - # Clean up filename output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_") - # Add output directory if specified if output_dir: output_path = os.path.join(output_dir, output_filename) else: output_path = output_filename - # Convert script = convert_notebook(content, source_name, allow_shell = allow_shell) - # Write output with open(output_path, "w", encoding = "utf-8") as f: f.write(script) @@ -349,11 +313,7 @@ Examples: ) parser.add_argument("notebooks", nargs = "+", help = "Notebook files or URLs to convert.") parser.add_argument("-o", "--output", dest = "output_dir", default = ".", help = "Output directory.") - # Default True for backwards compatibility: existing Colab notebooks - # routinely use pipes / redirection / interpolation in `!cmd` lines - # and the converted script needs to keep working. Operators who - # convert untrusted notebooks should pass --no-allow-shell to force - # a hard error on every metacharacter-bearing cell. + # Default True for backwards compat; pass --no-allow-shell for untrusted notebooks. parser.add_argument( "--allow-shell", dest = "allow_shell", @@ -371,14 +331,9 @@ Examples: args = parser.parse_args() - # Create output directory if needed os.makedirs(args.output_dir, exist_ok = True) - # SF2: track per-notebook failures so a CI invocation that converts - # 10 notebooks but silently fails on 3 is no longer reported as - # success. Each failure is collected and the loop continues so the - # caller sees the full set; final exit status is 1 if anything - # failed. + # Track per-notebook failures; continue the loop and exit 1 if any failed. failures: list[tuple[str, str]] = [] ok = 0 total = len(args.notebooks) diff --git a/scripts/notebook_validator.py b/scripts/notebook_validator.py index 74bfaa8d41..c1be7a63a4 100644 --- a/scripts/notebook_validator.py +++ b/scripts/notebook_validator.py @@ -49,12 +49,9 @@ from typing import Any, Iterable, Iterator def _atomic_write_bytes(path: pathlib.Path, data: bytes) -> None: - """Atomic write helper. See `scripts/scan_packages.py::update_req_file`. - - A crash between `mkstemp` and `os.replace` leaves the prior file - untouched, so a half-downloaded PyPI metadata cache file cannot - poison subsequent runs of the validator. - """ + """Atomic write (see scripts/scan_packages.py::update_req_file). A crash + between mkstemp and os.replace leaves the prior file intact, so a + half-downloaded cache file can't poison later runs.""" path.parent.mkdir(parents = True, exist_ok = True) dirpath = str(path.parent) or "." fd, tmp_path = tempfile.mkstemp(prefix = ".nb_val.", dir = dirpath) @@ -81,12 +78,10 @@ COLAB_PIP_FREEZE_URL = ( ) COLAB_FALLBACK_FILE = DATA_DIR / "colab_pip_freeze.gpu.txt" -# Oracle files we snapshot from googlecolab/backend-info. The diff -# subcommand fetches each, compares against the committed snapshot, -# and surfaces NEW / REMOVED / CHANGED entries so upstream Colab base -# image rotations land in CI within ~24h instead of when a notebook -# breaks. Every rule in this validator that resolves against the -# Colab preinstall (R-INST-002/003/004/005) gets earlier signal. +# Oracle files snapshotted from googlecolab/backend-info. The colab-diff +# subcommand surfaces NEW/REMOVED/CHANGED entries so upstream Colab base +# image rotations land in CI within ~24h, giving R-INST-002/003/004/005 +# earlier signal. COLAB_ORACLE_FILES: dict[str, str] = { "pip-freeze.gpu.txt": "colab_pip_freeze.gpu.txt", "apt-list-gpu.txt": "colab_apt_list.gpu.txt", @@ -145,9 +140,8 @@ class Finding: def iter_notebooks( notebooks_dir: pathlib.Path, include_templates: bool = False ) -> Iterator[pathlib.Path]: - """Yield user-facing .ipynb files under nb/ and kaggle/. Pass - include_templates=True to also walk original_template/ (used by the - convert subcommand which doesn't lint install cells).""" + """Yield user-facing .ipynb files under nb/ and kaggle/. + include_templates=True also walks original_template/ (for convert).""" subs = ("nb", "kaggle") if include_templates: subs = ("nb", "kaggle", "original_template") @@ -198,11 +192,8 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]: return out -# Notebook target environment. The Colab oracle (pip-freeze.gpu.txt) only -# applies to notebooks that actually run on Colab; AMD-Dev-Cloud, -# Kaggle, HuggingFace-Course, and DGX-Spark notebooks have their own -# preinstalled environments and the Colab-vs-cell rules are not -# applicable to them. +# Colab oracle only applies to notebooks that run on Colab; AMD, Kaggle, +# DGX-Spark have their own preinstalls and the Colab-vs-cell rules don't apply. def target_environment(notebook_name: str) -> str: parts = pathlib.PurePath(notebook_name).parts base = parts[-1] if parts else notebook_name @@ -466,19 +457,12 @@ def constraint_satisfied(version: str, ops: list[tuple[str, str]]) -> bool: def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]: - """Merge install-cell explicit constraints with Colab pip-freeze. Cell - wins. + """Merge install-cell constraints with Colab pip-freeze (cell wins). - Resolution order per package, when more than one form is present: - 1. Exact `==V` pin in any install line (definitive). - 2. Upper-bound `<=V` constraint (pip picks the highest - allowed; that's V). - 3. Colab pip-freeze fallback. - - The lower-bound `>=V` is intentionally NOT reflected here — a `>=V` - by itself doesn't change the resolved version when a higher - Colab-preinstalled version is already in scope. (R-INST-003 calls - `_install_cell_lower_bound` separately to model that case.) + Resolution order per package: (1) exact `==V` pin, (2) upper-bound `<=V` + (pip picks the highest allowed = V), (3) Colab fallback. Lower-bound `>=V` + is intentionally NOT reflected (it doesn't lower an already-higher Colab + version); R-INST-003 models that via `_install_cell_lower_bound`. """ out = dict(colab) pinned: set[str] = set() @@ -543,8 +527,7 @@ def rule_inst_002_no_deps_transitive( v = explicit_pin(sp) if v is None: continue - # Check transitive constraints on a curated short list of pkgs we - # care about (transformers/peft/trl/accelerate/torchao/torchcodec). + # Check transitive constraints on a curated short list of pkgs. for target in ( "tokenizers", "torchao", @@ -575,10 +558,9 @@ def rule_inst_002_no_deps_transitive( def _install_cell_lower_bound(install_cell: str, target: str) -> str | None: - """Return the highest LOWER bound that any install line places on `target`, - or None if no constraint is present. Treats `==V` as both lower and upper. - Used by R-INST-003: a `pip install torchao>=0.16.0` line is enough to - satisfy a `torchao>=0.16.0` floor even though it's not a `==` pin.""" + """Return the highest lower bound any install line places on `target` + (treating `==V` as both bounds), or None. Used by R-INST-003 so a + `torchao>=0.16.0` line satisfies the floor without a `==` pin.""" best: str | None = None for inv in iter_pip_invocations(install_cell): for raw in inv.packages: @@ -652,20 +634,17 @@ def rule_inst_004_torchcodec_torch( def rule_inst_005_transformers_tokenizers( install_cell: str, colab: dict[str, str], file: str, cell_idx: int ) -> list[Finding]: - """Fires only when transformers is installed with `--no-deps`. Without - `--no-deps`, pip resolves the correct tokenizers transitively, so the - rule would be a false positive (this is the case for older notebooks - that pin `transformers==4.51.3` but rely on pip's transitive resolver). - The rule targets the exact pattern PR #261b / #264 fixed: - `pip install --no-deps transformers==X` next to a Colab preinstall - `tokenizers` outside transformers's window.""" + """Fires only when transformers is installed with `--no-deps` (otherwise + pip resolves tokenizers transitively and flagging would be a false + positive). Targets the PR #261b/#264 pattern: `--no-deps transformers==X` + next to a Colab `tokenizers` outside transformers's window.""" findings: list[Finding] = [] res = resolved_set(install_cell, colab) tf = res.get("transformers") tok = res.get("tokenizers") if not tf or tok is None: return findings - # Find the install line that pins transformers and check for --no-deps. + # Find the transformers pin and check for --no-deps. transformers_line_no_deps = False for inv in iter_pip_invocations(install_cell): for raw in inv.packages: @@ -724,11 +703,9 @@ def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> li class _APIScanner(ast.NodeVisitor): """Scan user-facing code cells for known deprecated patterns. R-API-001 - (`for_training`/`for_inference`) is intentionally absent: those helpers - are still part of the live unsloth surface as of 2026-05; PR #221 removed - the calls cosmetically from Vision notebooks but did not deprecate the - methods. R-API-004 (live API surface diff) catches actual removals - dynamically without us hand-coding them.""" + (`for_training`/`for_inference`) is intentionally absent: those helpers are + still live as of 2026-05 (PR #221 removed them cosmetically, not as a + deprecation). R-API-004 catches actual removals dynamically.""" def __init__(self, file: str, cell_idx: int): self.file = file @@ -736,14 +713,10 @@ class _APIScanner(ast.NodeVisitor): self.findings: list[Finding] = [] def visit_Call(self, node: ast.Call) -> None: - # SFTConfig with suboptimal optim (R-API-003). - # NOTE: PR #221 also stripped `gradient_checkpointing` / - # `gradient_checkpointing_kwargs` from a handful of vision notebooks, - # but those kwargs are still accepted by live TRL (verified against - # trl==0.25.1 in the unsloth workspace) so removing them was - # cosmetic, not a deprecation. We do NOT flag them. R-API-004 (live - # API surface diff in the api subcommand) is the right way to catch - # actual TRL signature drift. + # SFTConfig with suboptimal optim (R-API-003). + # NOTE: PR #221 also stripped gradient_checkpointing kwargs from some + # vision notebooks, but they're still accepted by live TRL (trl==0.25.1) + # so that was cosmetic. We don't flag them; R-API-004 catches real drift. if isinstance(node.func, ast.Name) and node.func.id == "SFTConfig": for kw in node.keywords: if ( @@ -799,13 +772,9 @@ POLICY_CLAUSES_DEFAULT = [ def extract_policy_clauses(update_script: pathlib.Path) -> list[tuple[str, re.Pattern[str], Any]]: - """Best-effort: scan update_all_notebooks.py for canonical phrases used by - multiple templates. Falls back to POLICY_CLAUSES_DEFAULT. - - Today we use POLICY_CLAUSES_DEFAULT directly; the regex form is - intentionally permissive so a template-side reword (e.g. comment changes) - doesn't cause false positives. New clauses become 1-line PRs to this list. - """ + """Best-effort scan of update_all_notebooks.py for canonical phrases; + falls back to POLICY_CLAUSES_DEFAULT (which we use directly today). The + permissive regexes avoid false positives on template rewords.""" return list(POLICY_CLAUSES_DEFAULT) @@ -868,14 +837,9 @@ def cmd_drift(args: argparse.Namespace) -> int: check = False, capture_output = True, ) - # SF3: the restore MUST run even on SystemExit / KeyboardInterrupt / - # segfault-propagated exception, otherwise the user's working tree - # silently stays rolled back into the stash. A bare try/finally - # (NOT try/except/finally) preserves the original exception and - # still runs the cleanup. The pre-existing try/except around - # `subprocess.run` of the updater is folded inside the new outer - # try so its early returns still happen, but the stash pop is - # protected. + # The restore MUST run even on SystemExit/KeyboardInterrupt, else the + # working tree stays rolled back into the stash. A bare try/finally keeps + # the original exception while still running the cleanup (stash pop). findings: list[Finding] = [] rc: int try: @@ -920,8 +884,7 @@ def cmd_drift(args: argparse.Namespace) -> int: ) rc = 0 if not findings else 1 finally: - # Restore the working tree. Both commands MUST run regardless of - # how the try block exited (including SystemExit/KeyboardInterrupt). + # Restore the working tree (both commands run regardless of exit path). subprocess.run( ["git", "-C", str(nbdir), "checkout", "."], check = False, @@ -1004,19 +967,16 @@ def cmd_lint(args: argparse.Namespace) -> int: continue rel = str(path.relative_to(nbdir)) env = target_environment(rel) - # The Colab oracle is the source of truth ONLY for Colab notebooks. - # Other targets (amd / kaggle / dgx_spark) have their own runtime - # preinstall sets that aren't tracked here yet, so we apply the - # environment-agnostic rules and skip the Colab-specific ones. + # Colab oracle applies only to Colab notebooks; other targets get the + # environment-agnostic rules only (their preinstalls aren't tracked). oracle = colab if env == "colab" else {} cells = install_cells(nb) - # Per-cell rules: forbid-pattern checks scoped to a single line. + # Per-cell forbid-pattern checks. for idx, cell in cells: findings += rule_inst_001_git_plus(cell, rel, idx) findings += rule_inst_006_double_bang(cell, rel, idx) - # Whole-notebook rules: a notebook's install steps are sometimes split - # across multiple cells (initial install + post-install bumps). Merge - # all install cells before resolving compat against Colab. + # Whole-notebook rules: install steps may span multiple cells, so merge + # before resolving compat against Colab. merged = "\n".join(c for _, c in cells) if env == "colab" and merged: first_cell = cells[0][0] if cells else None @@ -1147,8 +1107,7 @@ def _parse_apt_lines(text: str) -> dict[str, str]: def _parse_os_lines(text: str) -> dict[str, str]: - """Free-form ` ` lines. Skip comments. The key is the - first token lower-cased; the value is the rest of the line.""" + """Free-form ` ` lines -> {tool_lower: rest}.""" out: dict[str, str] = {} for line in text.splitlines(): line = line.strip() @@ -1185,10 +1144,9 @@ def _diff_oracle( def cmd_colab_diff(args: argparse.Namespace) -> int: - """Fetch every Colab oracle file in COLAB_ORACLE_FILES, diff against - the committed snapshot, and print NEW / REMOVED / CHANGED. Advisory - by default (rc=0); --strict promotes any diff to rc=1 so the daily - cron can fail loudly when upstream rotates.""" + """Diff each Colab oracle file against its committed snapshot and print + NEW/REMOVED/CHANGED. Advisory (rc=0) by default; --strict makes any diff + rc=1 so the daily cron fails loudly on upstream rotation.""" snapshot_dir = pathlib.Path(args.snapshot_dir).resolve() any_diff = False for upstream_name, snapshot_name in COLAB_ORACLE_FILES.items(): diff --git a/scripts/run_ruff_format.py b/scripts/run_ruff_format.py index 2f0a6c6bba..817565eddf 100755 --- a/scripts/run_ruff_format.py +++ b/scripts/run_ruff_format.py @@ -19,10 +19,8 @@ def main(argv: list[str]) -> int: spacing_script = HERE / "enforce_kwargs_spacing.py" - # Pre-ruff: normalize def-signature trailing commas (>=3 params with a - # default -> one-per-line; everything else collapsible) and strip the magic - # comma from a short multi-line assert, so ruff wraps signatures accordingly - # and joins the assert back onto one line. + # Pre-ruff: normalize def-signature magic commas and strip the magic comma + # from short multi-line asserts so ruff wraps/joins accordingly. pre_cmd = [sys.executable, str(spacing_script), "--pre", *files] pre_proc = subprocess.run(pre_cmd) if pre_proc.returncode != 0: diff --git a/scripts/scan_npm_packages.py b/scripts/scan_npm_packages.py index 14a3cd6560..fe90afa7e6 100644 --- a/scripts/scan_npm_packages.py +++ b/scripts/scan_npm_packages.py @@ -7,73 +7,32 @@ """scan_npm_packages.py -- npm-side content scanner. -Counterpart to scripts/scan_packages.py for the pip ecosystem. Reads +npm counterpart to scripts/scan_packages.py. Reads studio/frontend/package-lock.json, downloads each resolved tarball DIRECTLY from registry.npmjs.org (never via `npm install` -- no -lifecycle scripts ever run), verifies the lockfile integrity hash, -unpacks each tarball into a sandboxed temp dir behind size / count / -path-escape / symlink guards, and pattern-scans the extracted file -contents for the signatures common to npm supply-chain attacks: +lifecycle scripts run), verifies the lockfile integrity hash, unpacks +each into a sandboxed temp dir behind size/count/path-escape/symlink +guards, and pattern-scans extracted contents for npm supply-chain +attack signatures: malicious lifecycle scripts, C2 / exfil hosts, +credential-stealing references, known IOC filenames, and obfuscation +shapes. - - Lifecycle (preinstall / install / postinstall / prepare) scripts - in any package.json that fetch + execute external code. - - C2 / exfiltration hosts (getsession.org, AWS IMDS endpoints, - Kubernetes ServiceAccount token paths, GitHub Actions OIDC, - HashiCorp Vault endpoints). - - Credential-stealing references (~/.npmrc, ~/.aws/credentials, - GITHUB_TOKEN / NPM_TOKEN in JS sources). - - Known IOC filenames from public advisories - (router_init.js, tanstack_runner.js, router_runtime.js). - - Obfuscation shapes (large single JS in package root with a low - whitespace ratio + Function/eval against a base64-decoded blob). - -Safety stance -============= - -This script ingests attacker-controlled archives. Every parse path -assumes the worst: - - 1. Downloads ONLY from `registry.npmjs.org`. Any tarball URL with a - different hostname is refused without fetching. - 2. Tarball download is size-capped (HARD_MAX_TARBALL_BYTES default - 64 MiB). HEAD-style probe via the Content-Length response header - plus a chunked read that aborts on overflow. +Safety stance (ingests attacker-controlled archives; assumes worst): + 1. Downloads ONLY from registry.npmjs.org; other hosts refused. + 2. Tarball download size-capped via Content-Length probe + chunked + read that aborts on overflow. 3. SHA-512 integrity verified against the lockfile entry BEFORE the - tarball is even opened. A mismatch aborts that package -- the - scanner does not "fall back" to the registry-published hash. - 4. tar extraction goes through `safe_extract`: - - rejects symbolic links (`SYMTYPE`, `LNKTYPE`) - - rejects absolute paths, `..` traversal, paths outside the - extract root after resolution - - rejects character / block / FIFO devices - - per-file uncompressed size cap (HARD_MAX_FILE_BYTES, default - 8 MiB) AND cumulative cap (HARD_MAX_TOTAL_BYTES, default - 128 MiB) AND member-count cap (HARD_MAX_MEMBERS, default - 50_000) - - tar reads happen via `tarfile.open(mode='r|gz')` streaming - so an oversized file is detected before write - 5. NOTHING from the extracted tree is ever executed. Files are read - as raw bytes, decoded with `errors='replace'`, and grepped. We - never call `node`, `eval`, `compile`, `subprocess.run`, - `os.system`, or anything that would touch the tarball's - declared scripts. - 6. Tempdir is created with `tempfile.mkdtemp(prefix='npm-scan-')`, - fully resolved with .resolve(), and registered with atexit to be - wiped on every termination path. - 7. Stdlib only. No third-party deps -- adding one would itself be a - supply-chain liability. + tarball is opened; mismatch aborts that package (no fallback). + 4. tar extraction via `safe_extract`: rejects symlinks, absolute / + `..` paths, device files; enforces per-file, cumulative, and + member-count caps; streams (`r|gz`) so oversize is caught early. + 5. NOTHING extracted is executed -- files are read as bytes and + grepped only. + 6. Tempdir resolved and atexit-wiped on every termination path. + 7. Stdlib only (a dep would be a supply-chain liability itself). -Exit codes -========== - - 0 no findings of severity HIGH or higher - 1 one or more HIGH/CRITICAL findings (or pre-scan structural - anomalies -- non-registry resolved URL, missing integrity) - 2 internal error (lockfile missing, integrity mismatch on - download, malformed tarball, etc.) - -The script is meant to be run in CI on every PR that touches -package-lock.json and on a nightly schedule. +Exit codes: 0 = no HIGH+ findings; 1 = HIGH/CRITICAL or pre-scan +structural anomaly; 2 = internal error. Run in CI per-PR and nightly. """ from __future__ import annotations @@ -565,12 +524,9 @@ CRED_HOST_NEEDS_CONTEXT: tuple[tuple[str, str], ...] = ( ), ) -# Credentials a frontend package should NEVER need to read. Bare -# substring match is too noisy (object-treeify ships a `docker` dev -# script that mounts ~/.npmrc -- legitimate dev tooling, never run -# at install time). We instead surface these only when they appear -# inside a LIFECYCLE script (preinstall / install / postinstall / -# prepare), which is the only path that runs automatically on +# Credentials a frontend package should never read. Bare substring +# match is too noisy (legit dev tooling mounts ~/.npmrc), so we flag +# these only inside lifecycle scripts -- the only auto-run path on # `npm ci`. See `scan_package_json` below. CRED_PATH_SUBSTRINGS: tuple[tuple[str, str], ...] = ( ("/.npmrc", "npm credentials file"), @@ -603,9 +559,8 @@ _JS_FETCH_EVAL = re.compile( """, ) -# `process.env.GITHUB_TOKEN` / `NPM_TOKEN` / `AWS_*` access in -# top-level / install-time code is suspicious. We also catch -# `os.environ["GITHUB_TOKEN"]` for the rare Python-in-npm postinstall. +# Token env access in install-time code; also catches os.environ[...] +# for the rare Python-in-npm postinstall. _JS_ENV_TOKEN = re.compile( r"""(process\.env\.|os\.environ\[?['"])(?: GITHUB_TOKEN | GH_TOKEN | NPM_TOKEN | NODE_AUTH_TOKEN @@ -616,11 +571,9 @@ _JS_ENV_TOKEN = re.compile( re.VERBOSE, ) -# Suspicious lifecycle-script payloads. Anything in a package.json -# `scripts` field that wgets/curls an external resource and executes -# it. We do NOT block ALL curl/wget in scripts (some legit packages -# fetch test fixtures into devDependencies), but we DO block the -# fetch+exec chain. +# Lifecycle-script fetch+exec chain: curl/wget an external resource +# and run it. Bare curl/wget is allowed (legit fixture fetches); only +# the fetch+exec chain is blocked. _LIFECYCLE_FETCH_EXEC = re.compile( r"""(?xs) (?:curl|wget|fetch|http\.get|axios\.get)\s+ # fetch verb @@ -634,9 +587,8 @@ _LIFECYCLE_FETCH_EXEC = re.compile( """, ) -# Obfuscation: large JS file that is mostly one line of base64-ish -# blob with a Function() / eval() bookend. Tuned against the -# router_init.js shape (2.3 MB obfuscated single-blob). +# Obfuscation: large single-line base64-ish blob behind Function()/ +# eval(). Tuned against the router_init.js shape (2.3 MB blob). _OBFUSC_BLOB = re.compile( r"""(?xs) (?:Function|eval)\s*\(\s*['"`]? @@ -653,11 +605,9 @@ _OBFUSC_BLOB = re.compile( def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]: """Return (entries, structural_findings). - Structural findings here are HIGH-severity refusals that should - short-circuit the scan -- a lockfile with non-registry resolved - URLs is itself a finding (covered by scripts/lockfile_supply_chain - _audit.py in detail; we surface a summary here so this scanner is - standalone-runnable). + Structural findings are HIGH-severity refusals that short-circuit + the scan (e.g. non-registry resolved URLs). A summary is surfaced + here so this scanner is standalone-runnable. """ entries: list[PackageEntry] = [] findings: list[Finding] = [] @@ -701,9 +651,8 @@ def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]: resolved = entry.get("resolved") if not resolved: continue - # Strict registry origin check. lockfile_supply_chain_audit - # already catches this; double-defend here so this scanner - # cannot be tricked into fetching from an attacker-chosen URL. + # Strict registry origin check so this scanner can't be tricked + # into fetching from an attacker-chosen URL. parsed = urllib.parse.urlparse(resolved) if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: findings.append( @@ -775,16 +724,12 @@ def download_tarball( timeout: float = HARD_HTTP_TIMEOUT_S, max_bytes: int = HARD_MAX_TARBALL_BYTES, ) -> tuple[Path, str | None]: - """Stream-download entry.resolved to dest. Verify SRI integrity. + """Stream-download entry.resolved to dest and verify SRI integrity. - Returns (downloaded_path, error_or_none). On any error the - returned path may not exist. Network access is restricted to - https://{ALLOWED_DOWNLOAD_HOST}/ -- the caller passes a Request - we already validated. + Returns (downloaded_path, error_or_none); on error the path may not + exist. Network access is restricted to ALLOWED_DOWNLOAD_HOST. """ - # Re-assert hostname; the entry was validated at parse time but a - # defence-in-depth check here means a future refactor cannot - # accidentally bypass it. + # Re-assert hostname (defence-in-depth against a future refactor). parsed = urllib.parse.urlparse(entry.resolved) if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST: return dest, (f"refused download from non-allowlisted URL {entry.resolved!r}") @@ -873,8 +818,7 @@ def safe_extract( total = 0 count = 0 try: - # Open in streaming mode so we never seek backwards in the - # input. `r|gz` rejects malformed gzip frames immediately. + # Streaming mode (no backward seeks); `r|gz` rejects bad gzip. with tarfile.open(tarball_path, mode = "r|gz") as tf: for member in tf: count += 1 @@ -889,9 +833,8 @@ def safe_extract( return f"refused link member {name!r} (sym/lnk)" if member.isdev() or member.isfifo(): return f"refused special member {name!r}" - # Cumulative cap is checked against DECLARED size up - # front to short-circuit obvious bombs without reading - # the body. + # Check declared size up front to short-circuit bombs + # without reading the body. declared = max(member.size, 0) if declared > HARD_MAX_BINARY_FILE_BYTES: return ( @@ -903,9 +846,8 @@ def safe_extract( f"cumulative bytes {total + declared} > cap " f"{max_total_bytes} at {name!r}" ) - # Strip leading "package/" -- the npm convention. We do - # NOT trust npm to be right, so we explicitly resolve - # the destination and refuse anything that escapes. + # Resolve destination and refuse anything escaping root + # (don't trust the npm "package/" convention). dest = extract_root / name if not _is_within(extract_root, dest): return f"refused escape: {name!r} resolved outside root" @@ -919,10 +861,8 @@ def safe_extract( src = tf.extractfile(member) if src is None: continue - # Sniff first 16 bytes to classify text vs binary. - # Text-cap members get the tight 16 MiB limit; binary - # members (executables, .node, .wasm, native libs) - # get the generous binary cap. We bound BOTH cases. + # Sniff first 16 bytes to classify text vs binary; + # each gets its own cap (both are bounded). header = src.read(16) is_binary = _looks_binary(name, header) file_cap = HARD_MAX_BINARY_FILE_BYTES if is_binary else HARD_MAX_TEXT_FILE_BYTES @@ -941,8 +881,7 @@ def safe_extract( f"({'binary' if is_binary else 'text'})" ) total += len(data) - # Write with restrictive mode (rw-r--r--) so even if - # someone runs the extract dir nothing is executable. + # Restrictive mode (rw-r--r--): nothing executable. with open(dest, "wb") as out: out.write(data) os.chmod(dest, 0o644) @@ -1008,10 +947,8 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ), ) ) - # Credential file paths inside a lifecycle script are - # exfiltration prep -- npm runs these scripts automatically - # on `npm ci`. Manual `scripts.*` entries (like a `docker` - # dev script) are out of scope: npm does not run them. + # Cred file paths in a lifecycle script are exfil prep (npm + # auto-runs these on `npm ci`); manual scripts are out of scope. for path_substr, why in CRED_PATH_SUBSTRINGS: if path_substr in body: findings.append( @@ -1071,20 +1008,12 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: def _host_in_outbound_context(text: str, host: str) -> bool: - """True if `host` appears in a way consistent with an outbound call. + """True if `host` appears consistent with an outbound call. - A bare `"169.254.169.254"` array literal (defensive blocklist) is - safe; a `fetch("http://169.254.169.254/...")` is not. The signal - is co-occurrence with either an HTTP URL scheme or a fetch verb - within a short window. - - A defensive blocklist looks like: - const CLOUD_METADATA_IPS = ["169.254.169.254", "169.254.170.2"]; - An exfil call looks like: - fetch("http://169.254.169.254/latest/meta-data/...") - http.request({ host: "169.254.169.254", path: "/..." }) + A bare array literal (defensive blocklist) is safe; co-occurrence + with an HTTP URL scheme or a fetch verb in a short window is not. """ - # Esc for use in a regex (IPs contain dots). + # Escape for regex (IPs contain dots). host_re = re.escape(host) # 1. URL form: http://host or https://host or //host/ or //host" url_form = re.compile( @@ -1127,8 +1056,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential surfaces. Tier 1: hosts with no legitimate use, - # bare substring is enough. + # Cred surfaces, tier 1: hosts with no legit use; bare substring. for needle, why in CRED_HOST_ALWAYS_BAD: if needle in text: findings.append( @@ -1145,8 +1073,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential surfaces. Tier 2: hosts that do appear in defensive - # code; require co-occurrence with a fetch verb or URL prefix. + # Cred surfaces, tier 2: hosts that appear in defensive code too; + # require co-occurrence with a fetch verb or URL prefix. for needle, why in CRED_HOST_NEEDS_CONTEXT: if needle in text and _host_in_outbound_context(text, needle): findings.append( @@ -1164,11 +1092,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: ) ) - # Credential PATHS are deliberately not scanned here; they have - # too high a false-positive rate at file scope (defensive code, - # docker mounts, AWS SDK docs strings). `scan_package_json` - # catches the malicious case -- credential paths inside a - # lifecycle script run automatically on `npm ci`. + # Credential PATHS aren't scanned here (too many FPs at file + # scope); scan_package_json catches them inside lifecycle scripts. # JS-specific regex. if _JS_FETCH_EVAL.search(text): @@ -1211,9 +1136,8 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]: return findings -# Filename suffix decides which scanners run. We deliberately treat -# *.cjs/*.mjs/*.ts the same as *.js -- attackers use whichever -# extension the consumer's bundler / loader resolves. +# Filename suffix decides which scanners run; .cjs/.mjs/.ts are +# treated like .js (attackers use whichever the loader resolves). _TEXT_SUFFIXES = ( ".js", ".mjs", @@ -1241,12 +1165,9 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: rel = path.relative_to(root).as_posix() lower = rel.lower() if not lower.endswith(_TEXT_SUFFIXES): - # Skip native binaries entirely -- regex over compiled - # machine code is just noise (false positives in WASM - # opcodes, .node BSS segments, image pixel data). Use - # content-magic detection so extensionless executables - # (eg `package/biome`) and versioned shared libraries - # are also skipped. + # Skip native binaries (regex over machine code is noise); + # content-magic detection also skips extensionless + # executables and versioned shared libraries. try: if path.stat().st_size > HARD_MAX_TEXT_FILE_BYTES: continue @@ -1288,12 +1209,12 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]: def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | None]: - """Download + extract + scan a single package. Cleans up its dir. + """Download + extract + scan a single package; cleans up its dir. Returns (findings, error). `error` is non-None only on hard - failures (download error, integrity mismatch, malformed tarball); - on a clean run with findings the error is None and the caller - decides exit code based on severity. + failures (download, integrity mismatch, malformed tarball); on a + clean run with findings, error is None and the caller decides the + exit code from severity. """ pkg_dir = workspace / f"{pkg.name.replace('/', '_')}-{pkg.version}" pkg_dir.mkdir(parents = True, exist_ok = True) diff --git a/scripts/scan_packages.py b/scripts/scan_packages.py index d753a17e49..861b35617b 100644 --- a/scripts/scan_packages.py +++ b/scripts/scan_packages.py @@ -56,27 +56,22 @@ from dataclasses import dataclass, field from pathlib import Path -# --------------------------------------------------------------------------- # Severity -# --------------------------------------------------------------------------- CRITICAL = "CRITICAL" HIGH = "HIGH" MEDIUM = "MEDIUM" SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2} -# Hard pin-blocks for publicly confirmed malicious PyPI versions. -# Source: Socket.dev 2026-05-12 disclosure (Mini Shai-Hulud May-12 wave) and -# earlier Semgrep / Endor reports for the `lightning` entries. +# Hard pin-blocks for confirmed malicious PyPI versions (Socket.dev 2026-05-12 +# Mini Shai-Hulud wave; earlier Semgrep/Endor reports for `lightning`). BLOCKED_PYPI_VERSIONS: dict[str, set[str]] = { "guardrails-ai": {"0.10.1"}, "mistralai": {"2.4.6"}, "lightning": {"2.6.2", "2.6.3"}, } -# --------------------------------------------------------------------------- # Pattern definitions -# --------------------------------------------------------------------------- # Subprocess / OS exec patterns RE_SUBPROCESS = re.compile( @@ -312,10 +307,8 @@ RE_C2_POLLING = re.compile( re.DOTALL, ) -# Developer-tool persistence hooks. The PyTorch Lightning 2.6.x compromise -# planted SessionStart hooks into Claude Code, VS Code tasks, and Cursor -# settings so the payload re-attached on every editor open. Catches any -# package writing into a known dev-tool config that supports auto-run. +# Developer-tool persistence hooks. Lightning 2.6.x planted SessionStart hooks +# into Claude Code / VS Code / Cursor so the payload re-attached on editor open. RE_DEV_TOOL_HIJACK = re.compile( r"\.claude/settings\.json" r"|\.cursor/.*hooks" @@ -326,9 +319,8 @@ RE_DEV_TOOL_HIJACK = re.compile( r"|\bautomator\b.*\.workflow\b", ) -# Hard-coded credential / API-token regexes embedded in source. Packages -# that ship regexes for OTHER people's secrets are nearly always -# stealers (litellm 1.82.7, elementary-data 0.23.3, Shai-Hulud). +# Hard-coded credential / API-token regexes embedded in source. Packages that +# ship regexes for OTHER people's secrets are nearly always stealers. RE_TOKEN_REGEX = re.compile( r"\bgh[psoru]_[A-Za-z0-9_]{20,}" # GitHub PAT/OAuth/etc. r"|\bgithub_pat_[A-Za-z0-9_]{20,}" @@ -342,20 +334,16 @@ RE_TOKEN_REGEX = re.compile( r"|\bglpat-[0-9A-Za-z_-]{20,}", # GitLab PAT ) -# Mini Shai-Hulud May-12 2026 wave indicators. The dropper artifact name -# `transformers.pyz` is high-confidence (no legit PyPI package ships a `.pyz` -# named after `transformers`); the host + slogans are CRITICAL. +# Mini Shai-Hulud May-12 2026 wave indicators. `transformers.pyz` dropper name +# is high-confidence; the host + slogans are CRITICAL. RE_MAY12_IOC = re.compile( r"(git-tanstack\.com|/tmp/transformers\.pyz|transformers\.pyz" r"|With Love TeamPCP|We've been online over 2 hours)", re.IGNORECASE, ) -# JavaScript-side obfuscation. The npm chalk/debug compromise and the -# Lightning router_runtime.js use the same minifier-style hex-var name -# pattern; a bundle full of `_0x1f2e3d` identifiers is a near-universal -# tell for a malicious npm payload (and very rare in legit minified code -# that ships in PyPI wheels). +# JavaScript-side obfuscation. A bundle full of `_0x1f2e3d` hex-var identifiers +# is a near-universal tell for a malicious npm payload, rare in legit wheels. RE_JS_OBFUSCATION = re.compile( r"_0x[a-f0-9]{4,6}\s*=\s*function" r"|var\s+_0x[a-f0-9]{4,6}\b" @@ -363,9 +351,8 @@ RE_JS_OBFUSCATION = re.compile( r"|String\.fromCharCode\s*\(\s*\d+\s*(?:,\s*\d+\s*){10,}\)", ) -# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch / -# XMLHttpRequest and attached a `window.ethereum` listener that -# Levenshtein-swapped recipient addresses on the way to the network. +# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch/XMLHttpRequest +# and swapped recipient addresses via a `window.ethereum` listener. RE_WEB3_HIJACK = re.compile( r"\bwindow\.ethereum\b" r"|\bweb3\.eth\.\w+\s*\(" @@ -374,11 +361,9 @@ RE_WEB3_HIJACK = re.compile( r"|TronWeb|solanaWeb3", ) -# Self-propagating supply-chain worms (Shai-Hulud, ForceMemo) plant -# their own GitHub workflow in every repo they can reach, and lean on -# trufflehog/gitleaks for credential discovery. The combo of any of -# these strings inside a *package payload* is overwhelming evidence of -# repo-takeover intent. +# Self-propagating worms (Shai-Hulud, ForceMemo) plant their own GitHub workflow +# in every repo they reach and use trufflehog/gitleaks for credential discovery. +# Any of these strings in a package payload is strong repo-takeover evidence. RE_WORKFLOW_INJECT = re.compile( r"\.github/workflows/[^\"\']*\.ya?ml" r"|\btrufflehog\b|\bgitleaks\b" @@ -388,9 +373,8 @@ RE_WORKFLOW_INJECT = re.compile( re.IGNORECASE | re.DOTALL, ) -# Shell-side patterns specific to install.sh / postinstall scripts that -# pipe remote code into a shell. `curl ... | sh` and friends are the -# canonical npm postinstall dropper. +# install.sh / postinstall scripts piping remote code into a shell. +# `curl ... | sh` is the canonical npm postinstall dropper. RE_SHELL_DROPPER = re.compile( r"\bcurl\b[^\n|]*\|\s*(?:sh|bash|zsh)\b" r"|\bwget\b[^\n|]*-O-\s*\|\s*(?:sh|bash|zsh)\b" @@ -400,9 +384,6 @@ RE_SHELL_DROPPER = re.compile( ) -# --------------------------------------------------------------------------- -# Finding dataclass -# --------------------------------------------------------------------------- @dataclass class Finding: severity: str @@ -412,9 +393,7 @@ class Finding: evidence: str = "" -# --------------------------------------------------------------------------- # Checkers -# --------------------------------------------------------------------------- def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: @@ -425,7 +404,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: """ findings = [] - # Only care about .pth files that have import lines (executable) + # Only .pth files with import lines are executable import_lines = [line for line in content.splitlines() if RE_PTH_IMPORT.match(line)] if not import_lines: return findings # Pure path entries, inert @@ -470,7 +449,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # Large base64 blob (special handling for blob size) + # Large base64 blob if RE_LARGE_BLOB.search(content): blob = RE_LARGE_BLOB.search(content).group() findings.append( @@ -483,7 +462,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # Catch-all: any import line at all in .pth (if nothing else triggered) + # Catch-all: any import line in .pth if nothing else triggered if not findings and import_lines: evidence = "\n".join(import_lines[:5]) if len(import_lines) > 5: @@ -521,7 +500,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: is_setup = basename in ("setup.py", "setup.cfg") is_init = basename == "__init__.py" - # Pre-compute all pattern matches + # Pre-compute pattern matches has_network = bool(RE_NETWORK.search(content)) has_subprocess = bool(RE_SUBPROCESS.search(content)) has_base64 = bool(RE_BASE64.search(content)) @@ -546,9 +525,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: has_c2_polling = bool(RE_C2_POLLING.search(content)) has_may12_ioc = bool(RE_MAY12_IOC.search(content)) - # --------------------------------------------------------------- # CRITICAL: combination patterns that strongly indicate malice - # --------------------------------------------------------------- # base64 decode + subprocess execution (staged payload) if has_base64 and has_subprocess: @@ -740,9 +717,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # --------------------------------------------------------------- # HIGH: single strong signals or weaker combinations - # --------------------------------------------------------------- # Obfuscated payload: base64 + exec/eval + large blob if has_base64 and has_exec_eval and has_blob: @@ -879,9 +854,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]: ) ) - # --------------------------------------------------------------- # MEDIUM: standalone signals (informational, may be legitimate) - # --------------------------------------------------------------- # base64 + exec/eval without blob if has_base64 and has_exec_eval and not has_blob: @@ -978,23 +951,18 @@ def _extract_evidence( return " | ".join(matches) if matches else "" -# --------------------------------------------------------------------------- # Non-Python checkers -# --------------------------------------------------------------------------- -# Several recent PyPI compromises (PyTorch Lightning 2.6.x, ForceMemo) -# carried the active payload in a bundled .js / .sh / workflow yaml so -# the Python imports looked clean on first glance. These checkers scan -# those file types when they appear inside a Python wheel/sdist. +# Recent PyPI compromises (Lightning 2.6.x, ForceMemo) carried the payload in a +# bundled .js / .sh / workflow yaml so the Python imports looked clean. These +# checkers scan those file types when they appear inside a wheel/sdist. def check_js_file(content: str, filename: str, package: str) -> list[Finding]: """Run JS-side checks. Triggered by .js / .mjs / .cjs / .ts.""" findings = [] - # A JS file *inside a Python wheel* that's larger than 100 KB is - # itself anomalous (legit Python packages don't ship hand-written - # JS bundles). Combined with ANY of the other JS heuristics it is - # CRITICAL; standalone it is HIGH. + # A >100 KB JS file inside a Python wheel is anomalous: CRITICAL combined + # with any other JS heuristic, HIGH standalone. is_large = len(content) > 100 * 1024 has_obf = bool(RE_JS_OBFUSCATION.search(content)) has_web3 = bool(RE_WEB3_HIJACK.search(content)) @@ -1119,10 +1087,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding] def check_workflow_file(content: str, filename: str, package: str) -> list[Finding]: """Run GitHub-Actions workflow checks. Triggered by .github/workflows/*.yml.""" findings = [] - # A GitHub workflow file inside a *PyPI package* is itself - # suspicious (Shai-Hulud's whole MO is to plant `shai-hulud.yml` - # in every repo it can write to). Anything matching the workflow - # injection signature gets flagged CRITICAL. + # A workflow file inside a PyPI package is suspicious (Shai-Hulud plants + # `shai-hulud.yml` everywhere); injection-signature matches are CRITICAL. if RE_WORKFLOW_INJECT.search(content): findings.append( Finding( @@ -1166,15 +1132,11 @@ def check_workflow_file(content: str, filename: str, package: str) -> list[Findi return findings -# --------------------------------------------------------------------------- # Archive handling -# --------------------------------------------------------------------------- # Tarbomb caps, mirrored from scripts/scan_npm_packages.py::safe_extract. -# Refuses zip-of-death / tar-of-death archives so a hostile sdist or -# wheel cannot exhaust memory or fill the temp dir before content -# scanning even starts. Keep these constants in sync with the npm side; -# we duplicate rather than import to keep `scan_packages.py` standalone. +# Refuses zip/tar-of-death so a hostile archive cannot exhaust memory before +# scanning. Keep in sync with the npm side; duplicated to stay standalone. HARD_MAX_FILE_BYTES = 64 * 1024 * 1024 # 64 MiB per member HARD_MAX_TOTAL_BYTES = 512 * 1024 * 1024 # 512 MiB cumulative HARD_MAX_MEMBERS = 50_000 # entries per archive @@ -1183,11 +1145,8 @@ HARD_MAX_MEMBERS = 50_000 # entries per archive def _refuse_unsafe_member_name(name: str) -> str | None: """Return a refusal reason for a member name, or None if safe. - Mirrors `scan_npm_packages.py::safe_extract` semantics: no absolute - paths, no `..` traversal segments. The caller is responsible for - checking the resolved path lands inside the extract root, but for - iter_archive_files we never write to disk so the name-shape check - plus the in-memory size cap is sufficient. + Mirrors `safe_extract`: no absolute paths, no `..` traversal. We never write + to disk, so the name-shape check plus the in-memory size cap is sufficient. """ if name.startswith("/") or ".." in Path(name).parts: return f"unsafe member name {name!r}" @@ -1197,9 +1156,8 @@ def _refuse_unsafe_member_name(name: str) -> str | None: def iter_archive_files(archive_path: str): """Yield (filename, text_content) for every file in a wheel/sdist. - Streams members with size + count caps applied at the member level - so a tarbomb / zipbomb cannot blow up the scanner's memory budget. - On cap breach we emit a `[WARN]` log and short-circuit the archive. + Streams members with per-member size + count caps so a tarbomb/zipbomb can't + blow the memory budget. On cap breach, emits a `[WARN]` and short-circuits. """ path = Path(archive_path) @@ -1225,7 +1183,7 @@ def iter_archive_files(archive_path: str): file = sys.stderr, ) continue - # Declared (uncompressed) size cap. + # Declared (uncompressed) size cap if info.file_size > HARD_MAX_FILE_BYTES: print( f" [WARN] {path.name}: skipped {info.filename!r} " @@ -1262,9 +1220,8 @@ def iter_archive_files(archive_path: str): file = sys.stderr, ) return - # Refuse symlinks / hardlinks / devices outright -- the - # scanner never writes them anyway, but tar parsers - # have historically dereferenced them on extract. + # Refuse symlinks/hardlinks/devices: tar parsers have + # historically dereferenced them on extract. if member.issym() or member.islnk(): print( f" [WARN] {path.name}: refused link member " f"{member.name!r}", @@ -1305,8 +1262,7 @@ def iter_archive_files(archive_path: str): f = tf.extractfile(member) if f is None: continue - # Bound the read so a tar header that lies about - # size cannot OOM us. + # Bound the read: a tar header may lie about size data = f.read(HARD_MAX_FILE_BYTES + 1) if len(data) > HARD_MAX_FILE_BYTES: print( @@ -1327,11 +1283,9 @@ def iter_archive_files(archive_path: str): def scan_archive(archive_path: str, package: str) -> list[Finding]: """Scan all files in an archive for malicious patterns. - A corrupted archive container (truncated wheel, bad gzip header, - etc.) used to be silently skipped by an ``except Exception: continue`` - inside ``iter_archive_files``. Per the silent-failure hardening - (SF1) it now emits a CRITICAL ``archive_corrupted`` finding so the - main loop counts and surfaces it rather than reporting "0 findings". + A corrupted archive container (truncated wheel, bad gzip header, etc.) emits + a CRITICAL ``archive_corrupted`` finding rather than being silently skipped + and reported as "0 findings" (silent-failure hardening SF1). """ findings: list[Finding] = [] try: @@ -1342,22 +1296,17 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]: elif lower.endswith(".py"): findings.extend(check_py_file(content, filename, package)) elif lower.endswith((".js", ".mjs", ".cjs", ".ts")): - # Lightning 2.6.x hid its real payload in a 14.8 MB - # router_runtime.js inside a Python wheel. Without this - # branch we'd have only seen the small Python loader. + # Lightning 2.6.x hid its payload in a 14.8 MB router_runtime.js; + # without this branch we'd only see the small Python loader. findings.extend(check_js_file(content, filename, package)) elif lower.endswith((".sh", ".bash")): findings.extend(check_shell_file(content, filename, package)) elif "/.github/workflows/" in lower and lower.endswith((".yml", ".yaml")): - # Shai-Hulud / ForceMemo plant their own GHA workflow. - # A workflow file inside a *PyPI package* is on its own - # already a yellow flag; pattern-match the worm signatures. + # Shai-Hulud/ForceMemo plant their own GHA workflow findings.extend(check_workflow_file(content, filename, package)) except (zipfile.BadZipFile, tarfile.TarError, EOFError, OSError) as exc: - # The archive cannot be opened or is structurally broken. A - # benign wheel/sdist always opens; a malformed one is either a - # transport corruption (treat as scan failure) or a deliberate - # attempt to bypass scanners that swallow archive errors. + # Archive cannot be opened / is structurally broken: either transport + # corruption or a deliberate attempt to bypass error-swallowing scanners. findings.append( Finding( CRITICAL, @@ -1370,9 +1319,7 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]: return findings -# --------------------------------------------------------------------------- # Download packages -# --------------------------------------------------------------------------- _RE_PYPI_SPEC_VERSION = re.compile(r"==\s*([A-Za-z0-9_.\-+!]+)") @@ -1382,10 +1329,8 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find """Filter ``specs`` against ``BLOCKED_PYPI_VERSIONS``. Returns ``(safe_specs, findings)``. Each blocked spec emits a CRITICAL - ``Finding`` and is removed from the returned spec list so the caller - never fetches the malicious tarball. Specs without an ``==X.Y.Z`` pin - pass through unchanged -- pip will resolve them at download time and - the existing scanners will catch the payload via the IOC regexes. + ``Finding`` and is dropped so the malicious tarball is never fetched. Specs + without an ``==X.Y.Z`` pin pass through; the IOC regexes catch them later. """ safe: list[str] = [] findings: list[Finding] = [] @@ -1407,7 +1352,7 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find f"{name}=={version} is on the BLOCKED_PYPI_VERSIONS list", ) ) - # Drop the spec; do not download. + # Drop the spec; do not download continue safe.append(spec) return safe, findings @@ -1416,14 +1361,11 @@ def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Find def _pip_download_env() -> dict[str, str]: """Return a scrubbed environment for invoking `pip download`. - Hostile shells / CI configs can override the index with PIP_INDEX_URL, - PIP_EXTRA_INDEX_URL, or a user `pip.conf`. We strip every PIP_* - override and route the resolver explicitly at PyPI. PIP_CONFIG_FILE - is forced to /dev/null so a stray ~/.pip/pip.conf with an - extra-index-url cannot bypass the pin. + Strips every PIP_* override and forces the resolver at PyPI; PIP_CONFIG_FILE + is /dev/null so a stray pip.conf extra-index-url cannot bypass the pin. """ env = {**os.environ} - # Drop any user override. + # Drop any user override for key in [k for k in env if k.startswith("PIP_")]: env.pop(key, None) env["PIP_INDEX_URL"] = "https://pypi.org/simple" @@ -1433,10 +1375,8 @@ def _pip_download_env() -> dict[str, str]: return env -# Pip resolver flags shared by both download branches. Pinning the -# index URL on the CLI is belt + braces with the env scrub above. -# `--no-build-isolation` is deliberately NOT set; we never invoke -# setup.py at all because of `--only-binary :all:`. +# Pip resolver flags shared by both download branches. CLI index-URL pin is +# belt + braces with the env scrub; `--only-binary :all:` avoids running setup.py. _PIP_DOWNLOAD_PIN_FLAGS = [ "--index-url", "https://pypi.org/simple", @@ -1445,9 +1385,8 @@ _PIP_DOWNLOAD_PIN_FLAGS = [ ] -# Strip any character that could escape `dest` via `os.path.join`. This -# is the last line of defence before `pkg_dir = os.path.join(dest, ...)` -# so a spec like `../../etc/foo==1.0` cannot land outside the temp tree. +# Strip characters that could escape `dest` via `os.path.join`, so a spec like +# `../../etc/foo==1.0` cannot land outside the temp tree. _RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]") @@ -1459,27 +1398,21 @@ def download_packages( ) -> tuple[list[tuple[str, str]], list[str]]: """Download packages to dest using pip download. NEVER installs. - Returns ``(results, download_errors)`` where ``results`` is a list of - ``(spec_or_name, filepath)`` for every downloaded archive and - ``download_errors`` is a list of one-line transport-failure summaries. - A non-empty ``download_errors`` MUST cause the caller to exit non-zero - even if no findings were produced; a silent ``0 findings, scan - incomplete`` is the bug class this return-shape was widened to fix. + Returns ``(results, download_errors)``: ``results`` is ``(spec_or_name, + filepath)`` per archive; ``download_errors`` is one-line transport-failure + summaries. A non-empty ``download_errors`` MUST make the caller exit + non-zero so a partial scan can't masquerade as "0 findings, all clean". - When with_deps=True, downloads the full transitive dependency tree - in a single pip invocation (all archives land in one flat dir). - When with_deps=False (default), downloads each spec individually - with --no-deps. + with_deps=True downloads the full transitive tree in one pip call (flat dir); + with_deps=False (default) downloads each spec individually with --no-deps. """ results: list[tuple[str, str]] = [] download_errors: list[str] = [] env = _pip_download_env() if with_deps: - # Single pip download call for all specs + their transitive deps. - # `--only-binary :all:` refuses sdists so we never execute a - # setup.py just to learn dependency metadata; combined with the - # scrubbed env, pip is wired hard at pypi.org. + # Single pip download for all specs + transitive deps. `--only-binary + # :all:` refuses sdists so we never execute setup.py for metadata. os.makedirs(dest, exist_ok = True) cmd = [ sys.executable, @@ -1495,7 +1428,7 @@ def download_packages( cmd, capture_output = True, text = True, - timeout = 600, # transitive resolution can be slow + timeout = 600, # transitive resolution is slow env = env, ) if proc.returncode != 0: @@ -1511,14 +1444,12 @@ def download_packages( for fname in sorted(os.listdir(dest)): fpath = os.path.join(dest, fname) if os.path.isfile(fpath): - # Derive package name from filename pkg_name = fname.split("-")[0].replace("_", "-").lower() results.append((pkg_name, fpath)) else: for spec in specs: raw_name = _extract_pkg_name(spec) - # Sanitize before joining into `dest` so a hostile spec - # cannot path-traverse out of the destination directory. + # Sanitize before joining into `dest` to prevent path traversal safe_name = _RE_PKG_NAME_SANITIZE.sub("_", raw_name) or "_pkg" pkg_dir = os.path.join(dest, safe_name) os.makedirs(pkg_dir, exist_ok = True) @@ -1552,7 +1483,6 @@ def download_packages( download_errors.append(msg) continue - # Find downloaded file(s) for fname in os.listdir(pkg_dir): fpath = os.path.join(pkg_dir, fname) if os.path.isfile(fpath): @@ -1560,9 +1490,7 @@ def download_packages( return results, download_errors -# --------------------------------------------------------------------------- # Parse requirements files -# --------------------------------------------------------------------------- _RE_NAME = re.compile(r"^([A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?)") @@ -1591,7 +1519,7 @@ def parse_requirements(req_files: list[str]) -> list[dict]: if not line or line.startswith("#") or line.startswith("-"): continue is_git = line.startswith("git+") or "git+" in line.split("#")[0] - # Strip inline comments and environment markers for spec + # Strip inline comments and env markers spec = line.split("#")[0].strip() spec = spec.split(";")[0].strip() if not spec: @@ -1624,7 +1552,7 @@ def get_downloaded_version(archive_path: str) -> str | None: parts = basename[:-4].split("-") if len(parts) >= 2: return parts[1] - # Sdist: name-version.tar.gz / .tar.bz2 / .zip + # Sdist: name-version. for ext in (".tar.gz", ".tar.bz2", ".tar.xz", ".tar", ".zip"): if basename.endswith(ext): stem = basename[: -len(ext)] @@ -1634,9 +1562,7 @@ def get_downloaded_version(archive_path: str) -> str | None: return None -# --------------------------------------------------------------------------- # Display -# --------------------------------------------------------------------------- def severity_color(sev: str) -> str: @@ -1652,7 +1578,6 @@ def print_findings(findings: list[Finding]) -> None: print("\n All clean. No suspicious patterns found.") return - # Sort by severity findings.sort(key = lambda f: SEVERITY_ORDER.get(f.severity, 99)) print(f"\n {'=' * 72}") @@ -1682,9 +1607,7 @@ def print_findings(findings: list[Finding]) -> None: print(f" Summary: {', '.join(parts)}") -# --------------------------------------------------------------------------- # PyPI version queries and --fix logic -# --------------------------------------------------------------------------- def version_sort_key(v: str) -> tuple: @@ -1708,15 +1631,13 @@ def version_sort_key(v: str) -> tuple: base = v_clean[0] suffix = v[len(base) :] - # Parse numeric parts parts = [] for seg in base.split("."): try: parts.append(int(seg)) except ValueError: parts.append(0) - # Pad to at least 3 parts - while len(parts) < 3: + while len(parts) < 3: # pad to at least 3 parts parts.append(0) # Suffix ordering: dev < alpha < beta < rc < (none) < post @@ -1732,7 +1653,7 @@ def version_sort_key(v: str) -> tuple: elif suffix_lower.startswith("post"): suffix_rank = 1 else: - suffix_rank = 0 # stable release + suffix_rank = 0 # stable return (epoch, tuple(parts), suffix_rank, suffix) @@ -1772,11 +1693,10 @@ def find_safe_version( print(f" [WARN] No versions found on PyPI for {name}", file = sys.stderr) return None - # Find index of bad version try: bad_idx = versions.index(bad_ver) except ValueError: - # bad_ver might have been resolved to a different string; search by sort key + # bad_ver may resolve to a different string; search by sort key bad_key = version_sort_key(bad_ver) bad_idx = None for i, v in enumerate(versions): @@ -1820,7 +1740,6 @@ def find_safe_version( print(f" {ver} -- CRITICAL finding(s), skipping") break - # Clean up scan dir for this version shutil.rmtree(scan_dir, ignore_errors = True) if clean: @@ -1850,8 +1769,7 @@ def update_req_line(raw_line: str, safe_ver: str, old_ver: str | None) -> str: code_part, marker = code_part.split(";", 1) marker = ";" + marker - # Replace version specifier - # Match patterns like ==1.2.3, >=1.2, ~=1.0, <=2.0, !=1.1, or bare name + # Replace version specifier (==1.2.3, >=1.2, ~=1.0, !=1.1, or bare name) rewritten = re.sub( r"([A-Za-z0-9._-]+)\s*(?:[><=!~]=?[^;#,\s]*(?:\s*,\s*[><=!~]=?[^;#,\s]*)*)?", lambda m: f"{m.group(1)}=={safe_ver}", @@ -1870,12 +1788,8 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: updates: {line_num (1-indexed): new_line_text} - Writes atomically: stage in a sibling tmp file on the same - filesystem, fsync, then `os.replace` over the original. A SIGKILL - or power loss mid-write therefore either leaves the original - intact or leaves the fully new file -- never a half-written - requirements file (which would silently re-introduce a malicious - pin). + Writes atomically (sibling tmp file, fsync, os.replace) so a crash mid-write + never leaves a half-written file that re-introduces a malicious pin. """ with open(filepath) as f: lines = f.readlines() @@ -1883,8 +1797,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: for line_num, new_text in updates.items(): idx = line_num - 1 if 0 <= idx < len(lines): - # Preserve original line ending - ending = "\n" if lines[idx].endswith("\n") else "" + ending = "\n" if lines[idx].endswith("\n") else "" # preserve line ending lines[idx] = new_text + ending dirpath = os.path.dirname(os.path.abspath(filepath)) or "." @@ -1909,7 +1822,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None: def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> None: """Run the --fix flow: find safe versions, update requirements files.""" - # Map package names to their entries for source tracking + # Map package names to entries for source tracking pkg_entries: dict[str, list[dict]] = {} for e in entries: norm = e["name"].lower().replace("-", "_").replace(".", "_") @@ -1931,8 +1844,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N changes_summary.append(f" SKIP {pkg_name} (git URL)") continue - # Get the currently resolved version - # Try to extract from the spec (e.g. name==1.2.3) + # Resolved version: try to extract from the spec (name==1.2.3) current_ver = None for e in related: spec = e["spec"] @@ -1947,7 +1859,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N downloaded = download_packages([pkg_name], dl_dir) if downloaded: current_ver = get_downloaded_version(downloaded[0][1]) - # Delete resolution download immediately shutil.rmtree(dl_dir, ignore_errors = True) if not current_ver: @@ -1986,7 +1897,6 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N for filepath, updates in file_updates.items(): update_req_file(filepath, updates) - # Print summary print(f"\n {'=' * 72}") print(f" FIX SUMMARY") print(f" {'=' * 72}") @@ -1995,9 +1905,7 @@ def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> N print(f"\n Re-run without --fix to verify the scan is clean.") -# --------------------------------------------------------------------------- # Directory scanning -# --------------------------------------------------------------------------- def _find_requirements_files(root: str) -> list[str]: @@ -2014,7 +1922,7 @@ def _find_requirements_files(root: str) -> list[str]: skip_dirs = {"__pycache__", "node_modules", "venv", ".venv", "site-packages"} results = [] for dirpath, dirnames, filenames in os.walk(root): - # Skip hidden dirs and known non-requirement dirs + # Skip hidden and known non-requirement dirs dirnames[:] = [ d for d in dirnames @@ -2024,18 +1932,15 @@ def _find_requirements_files(root: str) -> list[str]: for fname in sorted(filenames): if not fname.endswith(".txt"): continue - # Match requirements*.txt anywhere if fnmatch.fnmatch(fname.lower(), "requirements*.txt"): results.append(os.path.join(dirpath, fname)) - # Match *.txt inside a directory named "requirements" + # *.txt inside a directory named "requirements" elif dirname == "requirements": results.append(os.path.join(dirpath, fname)) return sorted(results) -# --------------------------------------------------------------------------- # Main -# --------------------------------------------------------------------------- def main() -> int: @@ -2134,7 +2039,7 @@ def main() -> int: all_findings: list[Finding] = [] - # Hard pin-block: refuse to download known-malicious PyPI versions. + # Hard pin-block: refuse to download known-malicious PyPI versions specs, blocked_findings = _check_blocked_pypi_versions(specs) all_findings.extend(blocked_findings) @@ -2172,10 +2077,8 @@ def main() -> int: ) _run_fix(critical_pkgs, entries, args.max_search) - # Surface any pip-download failures BEFORE the scan-result exit code so - # an empty / partial download cannot mask itself as "0 findings, all - # clean". This is item (4) of the silent-failure hardening: an - # unresolvable spec or PyPI timeout used to print to stderr and exit 0. + # Surface pip-download failures BEFORE the exit code so a partial download + # can't masquerade as "0 findings, all clean" (silent-failure hardening 4). if download_errors: print( f"\n {'=' * 72}\n" diff --git a/scripts/stamp_studio_release.py b/scripts/stamp_studio_release.py index a3c8712a6e..7dab35ea8a 100644 --- a/scripts/stamp_studio_release.py +++ b/scripts/stamp_studio_release.py @@ -22,11 +22,8 @@ def _atomic_write_text( data: str, encoding: str = "utf-8", ) -> None: - """Atomic version of ``Path.write_text``. - - A crash or signal mid-write leaves the prior file intact; the - Studio build never reads a partial ``_studio_release_build.py``. - """ + """Atomic ``Path.write_text``: a crash mid-write leaves the prior file + intact, so the build never reads a partial ``_studio_release_build.py``.""" dirpath = str(path.parent) or "." path.parent.mkdir(parents = True, exist_ok = True) fd, tmp_path = tempfile.mkstemp(prefix = ".stamp_studio.", dir = dirpath) diff --git a/scripts/verify_comment_only_diff.py b/scripts/verify_comment_only_diff.py index b661f76859..23d06b85df 100644 --- a/scripts/verify_comment_only_diff.py +++ b/scripts/verify_comment_only_diff.py @@ -57,9 +57,7 @@ def _git_show(rev: str, path: str) -> str: def _strip_docstrings(tree: ast.AST) -> ast.AST: - """Remove every string-literal docstring (Module / FunctionDef / - AsyncFunctionDef / ClassDef). Empty body becomes ``pass`` so - ast.unparse stays valid.""" + """Remove docstrings; empty bodies become ``pass`` so unparse stays valid.""" for node in ast.walk(tree): if isinstance( node, @@ -87,9 +85,8 @@ def _normalize_py(src: str) -> str: def _strip_shell_comments(s: str) -> str: - """Strip pure-comment lines and inline trailing comments from a shell - snippet, then collapse runs of blank lines. Heuristic only: leaves a - line untouched if it has an odd quote count (open string).""" + """Strip shell comments and collapse blank lines. Heuristic: skips lines + with an odd quote count (open string).""" out = [] for line in s.splitlines(): stripped = line.lstrip() @@ -116,9 +113,7 @@ def _strip_shell_comments(s: str) -> str: def _normalize_yaml_run_strings(obj: Any) -> Any: - """Walk the parsed YAML object; for any multi-line string (i.e. a - ``run: |`` script body), strip shell comments. Returns a normalised - copy.""" + """Strip shell comments from any multi-line string (``run: |`` body).""" if isinstance(obj, dict): return {k: _normalize_yaml_run_strings(v) for k, v in obj.items()} if isinstance(obj, list): diff --git a/scripts/verify_import_hoist.py b/scripts/verify_import_hoist.py index 6f5b452f3c..b8cb0573fe 100644 --- a/scripts/verify_import_hoist.py +++ b/scripts/verify_import_hoist.py @@ -122,16 +122,13 @@ class _Builder(ast.NodeVisitor): def __init__(self): self.module = Scope("module", "", None) - self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads - self.soft_uses: list[tuple[Scope, str, int]] = [] # annotations: count as "used" - # but never as "unresolved" - # (forward refs / string annos) + self.uses: list[tuple[Scope, str, int]] = [] # hard loads + # annotations: count as "used" but never as "unresolved" (forward refs) + self.soft_uses: list[tuple[Scope, str, int]] = [] def _visit_annotation(self, node, scope: Scope) -> None: - """Annotation context: with `from __future__ import annotations` these are - never evaluated (strings), and even otherwise they routinely contain forward - references. Record contained names as SOFT uses so an import used only in an - annotation still counts as used, but a forward-ref name is never 'unresolved'.""" + """Record annotation names as SOFT uses: an import used only in an annotation + counts as used, but a forward-ref name is never 'unresolved'.""" if node is None: return for n in ast.walk(node): @@ -189,7 +186,7 @@ class _Builder(ast.NodeVisitor): child = Scope("function", f"{scope.qualname}.{node.name}", scope) self._bind_type_params(node, child) self._bind_args(node.args, child) - # arg + return annotations: soft uses (may be strings / forward refs) + # arg + return annotations: soft uses for a in self._all_args(node.args): self._visit_annotation(a.annotation, child) self._visit_annotation(getattr(node, "returns", None), child) @@ -355,7 +352,7 @@ class _Builder(ast.NodeVisitor): if isinstance(node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)): child = Scope("comp", f"{scope.qualname}.", scope) for i, gen in enumerate(node.generators): - # first iterable is evaluated in the enclosing scope + # first iterable evaluates in the enclosing scope self._visit_expr(gen.iter, scope if i == 0 else child) self._bind_targets(child, gen.target) for cond in gen.ifs: @@ -394,9 +391,8 @@ def _any_star(scope: Scope) -> bool: def _resolve(scope: Scope, name: str): - """LEGB resolution. Returns (status, bindings) where status in + """LEGB resolution. Returns (status, bindings); status in {'local','import','other','builtin','star','unresolved'}.""" - # global / nonlocal redirection start = scope if name in scope.globals: chain = [_module_of(scope)] @@ -441,7 +437,7 @@ def _legb_chain(scope: Scope) -> list[Scope]: chain = [scope] p = scope.parent while p is not None: - if p.kind != "class" or p.parent is None: # module-level class never happens; keep module + if p.kind != "class" or p.parent is None: # skip class scopes, keep module if p.kind != "class": chain.append(p) p = p.parent @@ -455,7 +451,7 @@ def _analyze(src: str): tree = ast.parse(src) b = _Builder() b.run(tree) - # Per-scope: unresolved load names, and import targets it resolves to. + # Per-scope: unresolved load names + import targets it resolves to. unresolved: dict[str, set[str]] = {} targets_by_scope: dict[str, set[str]] = {} target_by_use: dict[tuple[str, str], set[str]] = {} @@ -467,7 +463,7 @@ def _analyze(src: str): tids = {bd.target for bd in binds if bd.target} targets_by_scope.setdefault(scope.qualname, set()).update(tids) target_by_use.setdefault((scope.qualname, name), set()).update(tids) - # soft uses (annotations): only contribute to "used", never to "unresolved" + # soft uses (annotations): contribute to "used" only, never "unresolved" for scope, name, _ln in b.soft_uses: status, binds = _resolve(scope, name) if status == "import": @@ -495,7 +491,7 @@ def _analyze(src: str): x.kind not in ("import", "importfrom") for x in bs ): ambiguous.setdefault(scope.qualname, set()).add(n) - # scope tree isn't stored; rebuild via uses is hard. We approximate with module only. + # scope tree isn't stored; approximate with module only. walk_scopes(module) return { @@ -524,14 +520,9 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] Blocker signals (precise, no relocation false-positives): UNRESOLVED-NEW - a load became undefined (dangling alias / removed import). - NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by - NO load. A correct hoist always wires its new import to a - reference; if the alias was left un-normalized OR renamed - to the wrong name, the hoisted import ends up unused. This - single signal catches BOTH user-described failure modes and - does NOT fire for code merely relocated to another file - (that removes the import, it doesn't add an unused one). - TARGET-CHANGED - the same (scope, name) load resolves to a different import + NEW-UNUSED-HOIST - a module-level import added by this change is resolved by + NO load (un-normalized alias or wrong rename target). + TARGET-CHANGED - same (scope, name) load resolves to a different import target before vs after (a same-name re-point). """ a = _analyze(before_src) @@ -566,14 +557,13 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] ) ) - # 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal) - # A module-level import in AFTER that NO load resolves to, and which was - # either newly added by this change OR was actually used before. Excludes: - # - relocation (the import is REMOVED, so it's not in after at all) - # - stable pre-existing re-exports (unused before AND after, not newly added) + # 2. HOISTED-IMPORT-UNUSED (core botched-hoist / wrong-rename signal) + # A module-level import in AFTER that NO load resolves to, that was either + # newly added by this change OR actually used before. Excludes relocation + # (import removed) and stable pre-existing re-exports. for n, tids in b["module_import_targets"].items(): if tids & after_used: - continue # resolved by something -> fine + continue # resolved -> fine newly_added = bool(tids - before_module_targets) was_used_before = bool(tids & before_used) if newly_added or was_used_before: @@ -619,8 +609,7 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]] findings.append(("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}")) # 6. TARGET-MISSING (informational): a scope stopped resolving to an import - # target. Real bugs are already covered above; remaining cases are code - # relocated to another file (e.g. a moved helper). Shown for transparency. + # target. Real bugs are covered above; remaining cases are relocated code. for scope, tbefore in a["targets_by_scope"].items(): tafter = b["targets_by_scope"].get(scope, set()) for t in sorted(tbefore - tafter): @@ -667,13 +656,13 @@ _SELF_TESTS = { "BLOCKER", ), "local_var_clash": ( - # _b renamed to b, but b is a LOCAL variable in f -> import silently unused + # _b renamed to b, but b is a LOCAL var in f -> import silently unused "def f(b):\n import re as _b\n return _b.compile(b)\n", "import re\ndef f(b):\n return b.compile(b)\n", # 'b' is the param, not the module "BLOCKER", ), "substring_safe": ( - # correct _copy->copy rename while a config_copy var exists: NO false positive + # correct _copy->copy rename while config_copy var exists: NO false positive "def f(config):\n" " import copy as _copy\n" " config_copy = _copy.deepcopy(config)\n" @@ -728,10 +717,9 @@ def _pyflakes_undefined(path: str) -> set[str] | None: def audit_files(paths: list[str]) -> int: - """Single-version robustness audit. For every file: confirm the analyzer does - not crash, then cross-check its 'unresolved' names against pyflakes. Any name - the resolver flags that pyflakes does NOT call undefined is a tool FALSE - POSITIVE (a resolver gap to fix).""" + """Single-version robustness audit: confirm the analyzer doesn't crash, then + cross-check its 'unresolved' names against pyflakes. A name the resolver flags + that pyflakes accepts is a tool false positive.""" n_files = n_err = n_fp = n_syntax = 0 fp_detail: dict[str, set[str]] = {} err_detail: dict[str, str] = {} diff --git a/studio/backend/_platform_compat.py b/studio/backend/_platform_compat.py index 37bf5e3fb5..013667d5ca 100644 --- a/studio/backend/_platform_compat.py +++ b/studio/backend/_platform_compat.py @@ -5,16 +5,15 @@ Compatibility shim for Anaconda/conda-forge Python builds. Anaconda puts distributor metadata between pipes in sys.version, e.g. -'3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'. The hardcoded -regex in platform._sys_version() can't parse this and raises ValueError; -CPython closed it as "not planned" (cpython#102396). +'3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'. The regex in +platform._sys_version() can't parse this and raises ValueError (cpython#102396, +closed as "not planned"). We seed platform._sys_version_cache so the stdlib parser never sees the bad string, fixing the import chain: structlog -> rich.pretty -> attrs._compat -> platform.python_implementation() -Import before any library that may trigger that chain. Idempotent (no-op if the -cache is already seeded or there are no pipes). +Import before any library that may trigger that chain. Idempotent. """ import platform @@ -23,14 +22,13 @@ import sys def _seed_sys_version_cache() -> None: - """One-shot cache prime: parse a cleaned sys.version and seed the cache.""" + """Parse a cleaned sys.version and seed the cache once.""" raw = sys.version # Strip paired |...| segments (Anaconda, conda-forge metadata) cleaned = re.sub(r"\s*\|[^|]*\|\s*", " ", raw).strip() - # Format B: "ver (build) | label | (build_dup) \n[compiler]" leaves two - # consecutive (...) groups after pipe-strip; drop the second. + # Pipe-strip can leave two consecutive (...) groups; drop the second. cleaned = re.sub(r"(\([^)]*\))\s+\([^)]*\)", r"\1", cleaned) if "|" in cleaned: @@ -43,7 +41,6 @@ def _seed_sys_version_cache() -> None: if cleaned == raw: return # Nothing to fix - # Parse the cleaned string through the real stdlib parser try: result = platform._sys_version(cleaned) except ValueError: diff --git a/studio/backend/auth/authentication.py b/studio/backend/auth/authentication.py index 6509cb6825..9dd56489eb 100644 --- a/studio/backend/auth/authentication.py +++ b/studio/backend/auth/authentication.py @@ -156,15 +156,7 @@ async def get_current_subject_allow_password_change( async def _get_current_subject( credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool ) -> str: - """ - FastAPI dependency: validate the JWT and return the subject. - - Use as a dependency on protected routes, e.g.: - - @router.get("/secure") - async def secure_endpoint(current_subject: str = Depends(get_current_subject)): - ... - """ + """FastAPI dependency: validate the JWT and return the subject. Use on protected routes.""" token = credentials.credentials # --- API key path (sk-unsloth-...) --- diff --git a/studio/backend/auth/hashing.py b/studio/backend/auth/hashing.py index 873381db04..0a1b0c15be 100644 --- a/studio/backend/auth/hashing.py +++ b/studio/backend/auth/hashing.py @@ -23,7 +23,7 @@ def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]: "sha256", password.encode("utf-8"), salt.encode("utf-8"), - 100_000, # 100k iterations + 100_000, ) return salt, dk.hex() diff --git a/studio/backend/auth/storage.py b/studio/backend/auth/storage.py index 42e2de7450..fa5b985513 100644 --- a/studio/backend/auth/storage.py +++ b/studio/backend/auth/storage.py @@ -26,30 +26,29 @@ _bootstrap_password: Optional[str] = None def generate_bootstrap_password() -> str: """Generate a 4-word diceware passphrase and persist it to disk. - Written to ``_BOOTSTRAP_PW_PATH`` so it survives restarts (the DB only - stores the *hash*). Later calls / restarts return the persisted value. + Persisted (the DB stores only the hash) so it survives restarts; later + calls return the persisted value. """ global _bootstrap_password - # 1. Already cached in this process? + # Cached in this process? if _bootstrap_password is not None: return _bootstrap_password - # 2. Already persisted from a previous run? + # Persisted from a previous run? if _BOOTSTRAP_PW_PATH.is_file(): _bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip() if _bootstrap_password: return _bootstrap_password - # 3. First-ever startup — generate a fresh passphrase. + # First startup: generate a fresh passphrase. import diceware _bootstrap_password = diceware.get_passphrase( options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"]) ) - # Persist so the *same* passphrase is reused if the server restarts - # before the user changes the password. + # Persist so the same passphrase survives restarts until password change. ensure_dir(_BOOTSTRAP_PW_PATH.parent) _BOOTSTRAP_PW_PATH.write_text(_bootstrap_password) try: @@ -87,17 +86,10 @@ def clear_bootstrap_password() -> None: def _hash_token(token: str) -> str: """SHA-256 hash helper for refresh token storage. - Plain SHA-256 is intentional: refresh tokens are high-entropy random - strings from ``secrets.token_urlsafe(48)`` (384 bits), so a slow KDF - (Argon2 / bcrypt / PBKDF2) adds zero security (2^384 is unbruteforceable - regardless of hash speed) while costing tens of ms per refresh. See the - OWASP Password Storage Cheat Sheet on hashing high-entropy inputs. - - API keys use the separate ``_pbkdf2_api_key`` helper (PBKDF2-HMAC-SHA256 - with a persistent server-side salt) — not for crypto reasons (128-bit - random tokens don't need slow hashing) but because CodeQL's - ``py/weak-sensitive-data-hashing`` query mislabels them as passwords and - demands a KDF. + Plain SHA-256 is intentional: refresh tokens are 384-bit random strings, so + a slow KDF adds no security while costing per-refresh latency. API keys use + the separate ``_pbkdf2_api_key`` helper, only to satisfy CodeQL's + ``py/weak-sensitive-data-hashing`` query, not for crypto reasons. """ return hashlib.sha256(token.encode("utf-8")).hexdigest() @@ -182,9 +174,8 @@ _api_key_pbkdf2_salt_cache: Optional[bytes] = None def _get_or_create_api_key_pbkdf2_salt() -> bytes: """Return the persistent API-key PBKDF2 salt, generating it once if missing. - Stored as a hex-encoded 32-byte random value in ``app_secrets`` under key - ``"api_key_pbkdf2_salt"``. Regenerated only when the row is missing (fresh - install, or operator deleted it and accepts invalidating existing keys). + Hex-encoded 32-byte random value in ``app_secrets``. Regenerated only when + the row is missing (fresh install, or operator deleted it). """ global _api_key_pbkdf2_salt_cache if _api_key_pbkdf2_salt_cache is not None: @@ -226,18 +217,10 @@ _DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at" def _pbkdf2_api_key(raw_key: str) -> str: """PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt. - For API-key storage ONLY, not refresh tokens. Matches the PBKDF2 algorithm - + iteration count used by the password hasher in ``auth/hashing.py`` so - the codebase is consistent on its credential-storage KDF. - - The slow KDF here is *only* a CodeQL appeasement, not a crypto - requirement: API keys are random 128-bit tokens (``secrets.token_hex``), - so brute force against 2^128 is infeasible regardless of hash speed. - CodeQL's ``py/weak-sensitive-data-hashing`` query mislabels them as - "password" data and demands a KDF from its allowlist (Argon2 / scrypt / - bcrypt / PBKDF2); we use PBKDF2 per its recommendation page. The salt is - still loaded from ``app_secrets`` so dumping the ``api_keys`` table alone - can't derive hashes for candidate tokens without the salt row. + For API-key storage ONLY, not refresh tokens. The slow KDF is only to + appease CodeQL's ``py/weak-sensitive-data-hashing`` query, not a crypto + requirement (API keys are random 128-bit tokens). The salt lives in + ``app_secrets`` so dumping ``api_keys`` alone can't derive hashes. """ salt = _get_or_create_api_key_pbkdf2_salt() dk = hashlib.pbkdf2_hmac( @@ -502,7 +485,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]: token_hash = _hash_token(token) conn = get_connection() try: - # Clean up any expired tokens while we're here + # Opportunistically clean up expired tokens conn.execute( "DELETE FROM refresh_tokens WHERE expires_at < ?", (datetime.now(timezone.utc).isoformat(),), diff --git a/studio/backend/colab.py b/studio/backend/colab.py index 1f618a34ab..ba46c52a6a 100644 --- a/studio/backend/colab.py +++ b/studio/backend/colab.py @@ -8,8 +8,7 @@ Colab helpers for Unsloth Studio. Uses Colab's built-in proxy. from pathlib import Path import sys -# Anaconda/conda-forge Python: seed platform._sys_version_cache before any -# import that triggers attrs -> rich -> structlog -> platform crash. +# Seed platform._sys_version_cache before attrs->rich->structlog->platform crash on conda Python. # See: https://github.com/python/cpython/issues/102396 _backend_dir = str(Path(__file__).parent) if _backend_dir not in sys.path: @@ -41,7 +40,7 @@ def get_colab_url(port: int = 8888) -> str: for attempt in range(3): try: url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10) - # Valid Colab proxy URL starts with https:// and embeds the port. + # Valid proxy URL is https:// and embeds the port. if url and isinstance(url, str) and url.startswith("https://") and str(port) in url: return url.rstrip("/") except Exception as e: @@ -59,16 +58,13 @@ def get_colab_url(port: int = 8888) -> str: def show_link(port: int = 8888, *, _url: "str | None" = None): """Display a styled clickable link to the UI. - *_url* is an optional pre-fetched Colab proxy URL; when omitted, - ``get_colab_url(port)`` is called. Pass it from ``_show_and_embed`` to - avoid a second ``eval_js`` round-trip. + *_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip. """ from IPython.display import display, HTML url = _url if _url is not None else get_colab_url(port) - # Truncated display URL. try/except so an unexpected URL shape never - # prevents the link from rendering. + # Truncated display URL; try/except so an odd URL shape still renders the link. try: port_prefix = f"{port}-" idx = url.index(port_prefix) @@ -77,7 +73,7 @@ def show_link(port: int = 8888, *, _url: "str | None" = None): except (ValueError, IndexError): short_url = url - # Plain-text line so the URL is visible even if HTML display fails. + # Plain-text line so the URL shows even if HTML display fails. logger.info(f"🌐 Unsloth Studio URL: {url}") html = f""" @@ -120,10 +116,8 @@ def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool: def _show_and_embed(port: int): """Embed the Studio inline for *port* with a branded header bar. - Fetches the Colab proxy URL once (also registering the port with Colab's - reverse-proxy) then renders a header bar + full-height iframe as one HTML - block. Falls back to ``serve_kernel_port_as_iframe`` if - ``IPython.display.HTML`` is unavailable. + Fetches the proxy URL once (registering the port), then renders header bar + + iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable. """ url = get_colab_url(port) logger.info(f"🌐 Unsloth Studio URL: {url}") @@ -162,7 +156,7 @@ def _show_and_embed(port: int): """) ) except Exception: - # Fallback: Colab's built-in helper (less control, but always works) + # Fallback: Colab's built-in helper. try: from google.colab import output as colab_output colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%") @@ -182,9 +176,8 @@ def start(port: int = 8888): logger.info("🦥 Starting Unsloth Studio...") - # --- Fast path: Studio is already running (cell re-run) --- - # Re-launching would collide on the port or silently shift to a new one. - # Just re-show the link and iframe instead. + # Fast path: Studio already running (cell re-run). Re-launching would collide on + # the port, so just re-show the link and iframe. if _is_studio_healthy(port): logger.info(f" Studio is already running on port {port} — reusing existing server.") _show_and_embed(port) @@ -217,16 +210,14 @@ def start(port: int = 8888): logger.error(f"❌ Unsloth Studio failed to start: {exc}") return - # run_server auto-increments the port if the requested one is in use (e.g. - # Jupyter on 8888). Read back the bound port so the Colab proxy URL and - # iframe point at the right place. + # run_server auto-increments the port if in use; read back the bound port so the + # proxy URL and iframe point at the right place. actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port logger.info(f" Server started on port {actual_port}!") - # Poll health endpoint to confirm reachability before showing the link and - # registering the iframe — avoids the race where ready_event fires but the - # process hasn't finished binding. + # Poll health endpoint before showing the link — avoids the race where ready_event + # fires but the process hasn't finished binding. import urllib.request server_ready = False @@ -247,9 +238,8 @@ def start(port: int = 8888): _show_and_embed(actual_port) - # Keep kernel alive so the daemon server thread stays running. Handle - # KeyboardInterrupt cleanly so interrupting the cell gives a readable - # message, not a raw traceback. + # Keep kernel alive so the daemon server thread runs; handle KeyboardInterrupt + # cleanly so interrupting the cell gives a readable message. try: for _ in range(10000): time.sleep(300) diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index 5400ef0d35..30ecb27f81 100644 --- a/studio/backend/core/__init__.py +++ b/studio/backend/core/__init__.py @@ -12,8 +12,8 @@ torch) before the version-activation code runs. import sys from pathlib import Path -# Put the backend dir on sys.path so bare "from utils.*" imports work when core -# is imported as a package (e.g. CLI: "from studio.backend.core import ModelConfig"). +# Add backend dir to sys.path so bare "from utils.*" imports work when core +# is imported as a package. _backend_dir = str(Path(__file__).resolve().parent.parent) if _backend_dir not in sys.path: sys.path.insert(0, _backend_dir) diff --git a/studio/backend/core/_torchao_stub.py b/studio/backend/core/_torchao_stub.py index 4a62004bac..6336954bd5 100644 --- a/studio/backend/core/_torchao_stub.py +++ b/studio/backend/core/_torchao_stub.py @@ -3,17 +3,11 @@ """Shared torchao Windows-ROCm import stub. -torchao (pulled in by transformers.quantizers) imports -torch.distributed._functional_collectives at module level, which imports -distributed_c10d.py unconditionally — that file crashes on Windows ROCm because -torch._C._distributed_c10d (the RCCL backend) is absent. `import -torch.distributed` alone is safe (guarded by `if is_available()`); the crash -only comes via torchao's import chain, so stubbing torchao short-circuits it. -_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports. - -Previously duplicated inline in run_export_process() and run_training_process(); -now both worker subprocesses call the single -`install_torchao_windows_rocm_stub()` entrypoint before importing +torchao (pulled in by transformers.quantizers) imports distributed_c10d.py +unconditionally, which crashes on Windows ROCm because the RCCL backend +(torch._C._distributed_c10d) is absent. Stubbing torchao short-circuits its +import chain; _StubSubpackageFinder handles any depth of torchao.xxx.yyy. +Worker subprocesses call install_torchao_windows_rocm_stub() before importing transformers / unsloth_zoo. """ @@ -27,12 +21,9 @@ import importlib.machinery _STUB_SENTINEL = object() -# Metaclass for stub types so isinstance(x, StubClass) returns False instead -# of raising TypeError ("arg 2 must be a type"). peft/tuners/lora/torchao.py -# does: -# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor -# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor)) -# If those names resolve to stub modules rather than types, isinstance() raises. +# Metaclass for stub types so isinstance(x, StubClass) returns False instead of +# raising TypeError -- peft's lora/torchao.py does isinstance() against torchao +# types, which fails if those names resolve to stub modules rather than types. class _StubTypeMeta(type): def __instancecheck__(cls, instance): return False @@ -70,8 +61,7 @@ def _make_mod_stub(mod_name): ): if attr.startswith("__"): raise AttributeError(attr) - # Return a stub CLASS (not a module) so isinstance(x, attr) returns - # False instead of raising TypeError. + # Return a stub CLASS (not module) so isinstance() returns False, not TypeError. child = _make_stub_type(f"{_n}.{attr}") setattr(_m, attr, child) return child @@ -113,14 +103,12 @@ class _StubSubpackageFinder(importlib.abc.MetaPathFinder): def install_torchao_windows_rocm_stub() -> None: """Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash. - No-op on every other platform (incl. Windows CUDA — there torchao is real - and shadowing it would break torchao quantization paths). Must run before - any import of transformers / unsloth_zoo. Safe to call once per worker. + No-op elsewhere (incl. Windows CUDA, where torchao is real). Must run before + importing transformers / unsloth_zoo. Safe to call once per worker. """ - # Gate on the active torch runtime, not env-var presence -- HIP_PATH / - # ROCM_PATH stay set after installing the HIP SDK and reverting to a CUDA - # torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip - # but still encode "rocm" in torch.__version__, so accept either. + # Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH + # persist after reverting to a CUDA wheel. Some ROCm wheels lack + # torch.version.hip but still encode "rocm" in __version__, so accept either. _is_win32_rocm = False if sys.platform == "win32": try: @@ -133,8 +121,7 @@ def install_torchao_windows_rocm_stub() -> None: except Exception: pass if _is_win32_rocm: - # Register the finder only on Windows ROCm -- elsewhere no stub modules - # are seeded, so appending would be pure accumulation. + # Register the finder only on Windows ROCm. sys.meta_path.append(_StubSubpackageFinder()) # Seed torchao top-level + key submodules; the finder handles the rest. for _tao_name in ( diff --git a/studio/backend/core/data_recipe/__init__.py b/studio/backend/core/data_recipe/__init__.py index 239e406300..ebc2798375 100644 --- a/studio/backend/core/data_recipe/__init__.py +++ b/studio/backend/core/data_recipe/__init__.py @@ -1,9 +1,7 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Data Recipe core (DataDesigner wrapper + job runner). -""" +"""Data Recipe core (DataDesigner wrapper + job runner).""" from .jobs import JobManager, get_job_manager diff --git a/studio/backend/core/data_recipe/huggingface.py b/studio/backend/core/data_recipe/huggingface.py index 3012fed1ea..7a1219b2c3 100644 --- a/studio/backend/core/data_recipe/huggingface.py +++ b/studio/backend/core/data_recipe/huggingface.py @@ -95,8 +95,7 @@ def publish_recipe_dataset( tags = None, ) card.text = card.text.replace(_DATA_DESIGNER_FOOTER, _UNSLOTH_STUDIO_FOOTER) - # Data Designer drops the explicit token when pushing the dataset - # card, so push it ourselves to keep auth request-local. + # Data Designer drops the explicit token, so push the card ourselves to keep auth request-local. card.push_to_hub(repo_id, token = hf_token, repo_type = "dataset") client._upload_main_dataset_files( diff --git a/studio/backend/core/data_recipe/jobs/manager.py b/studio/backend/core/data_recipe/jobs/manager.py index 71c2b7468f..32523e469e 100644 --- a/studio/backend/core/data_recipe/jobs/manager.py +++ b/studio/backend/core/data_recipe/jobs/manager.py @@ -132,10 +132,9 @@ class JobManager: ) -> str: """Spawn the job subprocess (one at a time, no cap). - ``internal_api_key_id`` is the row id of a workflow-scoped - sk-unsloth-* key minted by the route layer for local providers. - JobManager revokes it on terminal state so the key's live window is - no longer than the run. + ``internal_api_key_id`` is a workflow-scoped sk-unsloth-* key row id + minted by the route layer; revoked on terminal state so the key's + live window is no longer than the run. """ llm_columns = recipe.get("columns") or [] llm_column_count = 0 @@ -537,15 +536,14 @@ class JobManager: def _retire_workflow_key(self, job: Job) -> None: """Revoke the workflow-scoped sk-unsloth-* key, if one was minted. - Best-effort: revocation failures are swallowed. The key expires on - its own after 24h, so a missed revoke is a latency concern, not a - correctness one. + Best-effort: failures are swallowed. The key expires after 24h, so a + missed revoke is a latency, not correctness, concern. """ key_id = getattr(job, "internal_api_key_id", None) if not key_id: return try: - from auth import storage # deferred: avoids circular import + from auth import storage # deferred: avoid circular import storage.revoke_internal_api_key(int(key_id)) except Exception: pass diff --git a/studio/backend/core/data_recipe/jobs/parse.py b/studio/backend/core/data_recipe/jobs/parse.py index e37d7064ed..3be830d0e4 100644 --- a/studio/backend/core/data_recipe/jobs/parse.py +++ b/studio/backend/core/data_recipe/jobs/parse.py @@ -327,7 +327,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None: _apply_source_progress(job, update.source_progress) if update.stage in USAGE_RESET_STAGES: - # Usage summary is a short block; reset on entering the next stage. + # Usage summary is a short block; reset on the next stage. job._in_usage_summary = False if update.usage_section_start is not None: diff --git a/studio/backend/core/data_recipe/jobs/types.py b/studio/backend/core/data_recipe/jobs/types.py index 6dd4db5973..c30606b2d1 100644 --- a/studio/backend/core/data_recipe/jobs/types.py +++ b/studio/backend/core/data_recipe/jobs/types.py @@ -90,9 +90,8 @@ class Job: progress_columns_total: int | None = None source_progress_estimated_total: int | None = None completed_columns: list[str] = field(default_factory = list) - # Id of the internal sk-unsloth-* API key minted for a local-model - # workflow. Revoked when the job ends so the key's window matches the - # run rather than its 24h TTL. + # Id of the internal sk-unsloth-* API key minted for a local-model workflow. + # Revoked when the job ends so the key's window matches the run, not its 24h TTL. internal_api_key_id: int | None = None _current_usage_model: str | None = None _in_usage_summary: bool = False diff --git a/studio/backend/core/data_recipe/jobs/worker.py b/studio/backend/core/data_recipe/jobs/worker.py index 27d08e90fe..4073288cb3 100644 --- a/studio/backend/core/data_recipe/jobs/worker.py +++ b/studio/backend/core/data_recipe/jobs/worker.py @@ -76,7 +76,7 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) """Subprocess entrypoint. Sends events to `event_queue`.""" import os - os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports import warnings from loggers.config import LogConfig @@ -112,8 +112,8 @@ def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) builder = build_config_builder(recipe) designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT)) - # DataDesigner configures root logging in __init__. Attach the queue - # logger to `data_designer` directly so parser events survive root resets. + # DataDesigner resets root logging in __init__; attach the queue handler + # to the named loggers directly so parser events survive. handler = _QueueLogHandler(event_queue) handler.setLevel(logging.INFO) for logger_name in ( diff --git a/studio/backend/core/data_recipe/local_callable_validators.py b/studio/backend/core/data_recipe/local_callable_validators.py index a40261fa92..fbe847f9ce 100644 --- a/studio/backend/core/data_recipe/local_callable_validators.py +++ b/studio/backend/core/data_recipe/local_callable_validators.py @@ -184,7 +184,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape: normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto" def _validator(df): - import pandas as pd # imported lazily for local callable runtime + import pandas as pd # lazy import for local callable runtime row_count = int(len(df.index)) if row_count == 0: diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index addbe8f848..9d8ca5cfcc 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -167,10 +167,8 @@ def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: li def build_mcp_providers(recipe: dict[str, Any]) -> list: from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports] - # Same gate as the chat MCP path: stdio providers spawn a local - # subprocess, so build them only when this host allows it (desktop / - # explicit opt-in). Otherwise a recipe carried onto a hosted host - # cannot spawn. + # Same gate as the chat MCP path: stdio providers spawn a local subprocess, + # so build them only when this host allows it (desktop / explicit opt-in). from core.inference.mcp_client import stdio_mcp_enabled stdio_allowed = stdio_mcp_enabled() @@ -283,9 +281,8 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = model_providers = build_model_providers(recipe) _validate_recipe_runtime_support(recipe, model_providers) - # DataDesigner requires at least one model provider in its registry even - # when the pipeline has no LLM columns. Supply a lightweight stub so - # sampler/expression-only recipes can run without a real provider. + # DataDesigner requires >=1 model provider even with no LLM columns; stub + # one so sampler/expression-only recipes run without a real provider. if not model_providers: from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports] model_providers = [ diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py index db9d1b352c..f7241cd925 100644 --- a/studio/backend/core/export/__init__.py +++ b/studio/backend/core/export/__init__.py @@ -1,12 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Export submodule - Model export operations +"""Export submodule - model export operations. -The default get_export_backend() returns an ExportOrchestrator that -delegates to a subprocess. The original ExportBackend runs inside the -subprocess and can be imported directly from .export when needed. +get_export_backend() returns an ExportOrchestrator that delegates to a +subprocess. The original ExportBackend runs inside the subprocess and can be +imported directly from .export when needed. """ from .orchestrator import ExportOrchestrator, get_export_backend diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index aaae8002fc..d5a3de06df 100644 --- a/studio/backend/core/export/export.py +++ b/studio/backend/core/export/export.py @@ -1,7 +1,6 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -# backend/export.py """Export backend - exports models in various formats.""" import glob @@ -44,10 +43,8 @@ def _is_wsl(): def _apply_wsl_sudo_patch(): """On WSL, monkey-patch do_we_need_sudo() to return False. - WSL lacks passwordless sudo, and do_we_need_sudo() runs - `sudo apt-get update`, which hangs on a stdin password in a - non-interactive subprocess. setup.sh pre-installs the build deps on - WSL, so sudo isn't needed at runtime. + WSL lacks passwordless sudo and do_we_need_sudo()'s `sudo apt-get update` + hangs on a stdin password; setup.sh pre-installs the build deps anyway. """ if not _is_wsl(): return @@ -108,18 +105,15 @@ class ExportBackend: try: logger.info("Starting memory cleanup...") - # Unload all inference-backend models model_names = list(self.inference_backend.models.keys()) for model_name in model_names: self.inference_backend.unload_model(model_name) - # Clear current export state self.current_model = None self.current_tokenizer = None self.current_checkpoint = None self._audio_type = None - # Clear GPU cache (handles gc + backend-specific cleanup) clear_gpu_cache() logger.info("Memory cleanup completed successfully") @@ -156,7 +150,6 @@ class ExportBackend: try: logger.info(f"Loading checkpoint: {checkpoint_path}") - # Cleanup existing models first self.cleanup_memory() checkpoint_path_obj = Path(checkpoint_path) @@ -171,11 +164,9 @@ class ExportBackend: model_id = base_model or checkpoint_path - # Detect audio type and vision self._audio_type = detect_audio_type(model_id) self.is_vision = not self._audio_type and is_vision_model(model_id) - # Load model based on type if self._audio_type == "csm": from unsloth import FastModel from transformers import CsmForConditionalGeneration @@ -255,7 +246,6 @@ class ExportBackend: trust_remote_code = trust_remote_code, ) - # Detect PEFT / LoRA model if _IS_MLX: # MLX doesn't use PeftModel — detect LoRA via adapter_config.json self.is_peft = adapter_config.exists() @@ -345,7 +335,6 @@ class ExportBackend: else: save_method = "merged_16bit" - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving merged model locally to: {save_directory}") @@ -366,7 +355,6 @@ class ExportBackend: logger.info(f"Model saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -447,7 +435,6 @@ class ExportBackend: output_path: Optional[str] = None try: - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving base model locally to: {save_directory}") @@ -470,7 +457,6 @@ class ExportBackend: logger.info(f"Model saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -519,7 +505,6 @@ class ExportBackend: ) username = repo_id.split("/")[0] - # Build and push model card content = MODEL_CARD.format( username = username, base_model = base_model, @@ -530,7 +515,6 @@ class ExportBackend: card = ModelCard(content) card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card") - # Upload files if save_directory: hf_api.upload_folder( folder_path = save_directory, @@ -583,10 +567,8 @@ class ExportBackend: # unsloth expects lowercase quant method quant_method = quantization_method.lower() - # Pin convert_hf_to_gguf.py to the same llama.cpp ref as the - # llama-quantize binary (Studio installs at a tagged ref via - # setup.sh) so it can't drift past the pinned binary's gguf API. - # Set before both branches; hub-only export has save_directory == "". + # Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it + # can't drift past the pinned llama-quantize binary's gguf API. global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED try: from unsloth_zoo.llama_cpp import ( @@ -605,12 +587,10 @@ class ExportBackend: ) _LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) - # Use absolute path so unsloth's relative-path internals - # (check_llama_cpp, use_local_gguf, _download_convert_hf_to_gguf) - # resolve against the repo root cwd, NOT the export directory. + # Absolute path so unsloth's relative-path internals resolve + # against the repo root cwd, not the export directory. abs_save_dir = os.path.abspath(save_directory) logger.info(f"Saving GGUF model locally to: {abs_save_dir}") @@ -619,15 +599,11 @@ class ExportBackend: # On WSL, patch out sudo check before llama.cpp build _apply_wsl_sudo_patch() - # Snapshot existing .gguf files in cwd before conversion; - # unsloth's convert_to_gguf writes output relative to cwd - # (repo root), so we diff afterwards and relocate. + # convert_to_gguf writes output relative to cwd (repo root); + # snapshot existing .gguf so we can diff and relocate afterwards. cwd = os.getcwd() pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - # Absolute path — no os.chdir needed. unsloth saves intermediate - # HF model files into model_save_path; unsloth-zoo's - # check_llama_cpp() uses ~/.unsloth/llama.cpp by default. model_save_path = os.path.join(abs_save_dir, "model") self.current_model.save_pretrained_gguf( model_save_path, @@ -635,18 +611,14 @@ class ExportBackend: quantization_method = quant_method, ) - # Relocate GGUF artifacts into the export directory. - # convert_to_gguf writes .gguf to cwd (repo root) because - # --outfile is a relative path like "model.Q4_K_M.gguf". + # Relocate the .gguf that convert_to_gguf wrote to cwd (repo root). new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs for src in sorted(new_ggufs): dest = os.path.join(abs_save_dir, os.path.basename(src)) shutil.move(src, dest) logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/") - # Flatten any .gguf from subdirs into abs_save_dir. - # save_pretrained_gguf may create subdirs (e.g. model_gguf/) - # named differently from model_save_path. + # Flatten any .gguf from subdirs (e.g. model_gguf/) into abs_save_dir. for sub in list(Path(abs_save_dir).iterdir()): if not sub.is_dir(): continue @@ -654,13 +626,11 @@ class ExportBackend: dest = os.path.join(abs_save_dir, src.name) shutil.move(str(src), dest) logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/") - # Clean up the subdir (intermediate HF files, etc.) shutil.rmtree(str(sub), ignore_errors = True) logger.info(f"Cleaned up subdirectory: {sub.name}") - # For non-PEFT models, save_pretrained_gguf redirects to the - # checkpoint path, leaving a *_gguf dir in outputs/. Relocate - # any GGUFs from there and clean it up. + # For non-PEFT models, save_pretrained_gguf leaves a *_gguf dir at + # the checkpoint path; relocate its GGUFs and clean it up. if self.current_checkpoint: ckpt = Path(self.current_checkpoint) gguf_dir = ckpt.parent / f"{ckpt.name}_gguf" @@ -680,7 +650,6 @@ class ExportBackend: # Write export metadata so the Chat page can identify the base model self._write_export_metadata(abs_save_dir) - # Log final file locations (post-relocation). final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf"))) logger.info( "GGUF export complete. Final files in %s:\n %s", @@ -689,7 +658,6 @@ class ExportBackend: ) output_path = str(Path(abs_save_dir).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( @@ -743,7 +711,6 @@ class ExportBackend: output_path: Optional[str] = None try: - # Save locally if requested if save_directory: save_directory = str(resolve_export_dir(save_directory)) logger.info(f"Saving LoRA adapter locally to: {save_directory}") @@ -759,7 +726,6 @@ class ExportBackend: logger.info(f"Adapter saved successfully to {save_directory}") output_path = str(Path(save_directory).resolve()) - # Push to hub if requested if push_to_hub: if not repo_id or not hf_token: return ( diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py index 49c1af8822..20158d1891 100644 --- a/studio/backend/core/export/orchestrator.py +++ b/studio/backend/core/export/orchestrator.py @@ -28,9 +28,7 @@ logger = get_logger(__name__) _CTX = mp.get_context("spawn") -# Max captured log lines kept in memory per export orchestrator; -# scrollback for the live export log panel. 4000 lines is ~1 MB -# worst-case at 256 chars/line. +# Max log lines kept per orchestrator (live log panel scrollback); ~1 MB worst-case. _LOG_BUFFER_MAXLEN = 4000 @@ -44,13 +42,10 @@ class ExportOrchestrator: """ def __init__(self): - # Subprocess state. self._proc: Optional[mp.Process] = None self._cmd_queue: Any = None self._resp_queue: Any = None - # Serializes export ops (load_checkpoint, export_*, cleanup) so - # concurrent HTTP requests can't interleave commands on the - # subprocess queue. + # Serializes export ops so concurrent HTTP requests can't interleave commands. self._lock = threading.Lock() # Local state mirrors (updated from subprocess responses). @@ -58,25 +53,15 @@ class ExportOrchestrator: self.is_vision: bool = False self.is_peft: bool = False - # ── Live log capture ───────────────────────────────────── - # Thread-safe ring buffer of log lines from the worker - # subprocess. Powers the GET /api/export/logs/stream SSE - # endpoint the export dialog consumes. + # Thread-safe ring buffer of worker log lines; powers the export logs SSE endpoint. self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN) self._log_lock = threading.Lock() - # Monotonic sequence number. Never reset across operations, so - # SSE clients can use it as a stable cursor even if clear_logs() - # runs mid-session. + # Monotonic seq, never reset, so SSE clients have a stable cursor across clear_logs(). self._log_seq: int = 0 - # Snapshot of _log_seq at the start of the current run (set by - # clear_logs()). The SSE endpoint defaults its cursor here so a - # client connecting AFTER the worker's first lines still sees the - # full run. Lines in the current run have seq > _run_start_seq; - # prior-run lines have seq <= it. + # _log_seq snapshot at the current run's start; SSE defaults its cursor here so a + # late-connecting client still sees the full run. Current run has seq > this. self._run_start_seq: int = 0 - # True while an export op (load/export/cleanup) is running. The - # SSE endpoint ends the stream 1s after this flips False to drain - # trailing log lines. + # True while an export op runs; SSE ends the stream 1s after this flips False. self._export_active: bool = False atexit.register(self._cleanup) @@ -87,12 +72,7 @@ class ExportOrchestrator: # ------------------------------------------------------------------ def _append_log(self, entry: Dict[str, Any]) -> None: - """Append a worker-subprocess log line to the buffer. - - Entries look like {"type": "log", "stream": "stdout"|"stderr", - "line": "...", "ts": ...}. Each gets a monotonic seq number so - SSE clients can cursor through new lines. - """ + """Append a worker log line to the buffer, stamped with a monotonic seq.""" line = entry.get("line") if not line: return @@ -108,17 +88,10 @@ class ExportOrchestrator: ) def clear_logs(self) -> None: - """Drop buffered log lines from a previous operation. + """Drop buffered log lines from a previous op so the UI shows only this run. - Called at the start of each export op so the UI shows only the - current run. The seq counter is NOT reset, so an SSE client that - captured the cursor before clear_logs() still sees new lines - (with strictly greater seq). - - Also snapshots the current seq into ``_run_start_seq`` so the SSE - endpoint can anchor its default cursor at this run's start. - Anything appended after has seq > the snapshot, reachable via - ``get_logs_since(get_run_start_seq())``. + The seq counter is NOT reset (clients keep a stable cursor); the current seq + is snapshotted into ``_run_start_seq`` to anchor the SSE default cursor. """ with self._log_lock: self._log_buffer.clear() @@ -138,12 +111,7 @@ class ExportOrchestrator: return self._log_seq def get_run_start_seq(self) -> int: - """Return the seq captured at the start of the current run. - - The SSE endpoint uses this as the default cursor so a client - connecting AFTER the worker started emitting still sees every - line from the current run. - """ + """Return the seq captured at the current run's start (SSE default cursor).""" with self._log_lock: return self._run_start_seq @@ -187,22 +155,19 @@ class ExportOrchestrator: self._proc = None return - # 1. Drain stale responses. self._drain_queue() - # 2. Send shutdown command. try: self._cmd_queue.put({"type": "shutdown"}) except (OSError, ValueError): pass - # 3. Wait for graceful shutdown. try: self._proc.join(timeout = timeout) except Exception: pass - # 4. Force kill if still alive. + # Force kill if still alive. if self._proc is not None and self._proc.is_alive(): logger.warning("Export subprocess did not exit gracefully, terminating") try: @@ -272,7 +237,6 @@ class ExportOrchestrator: resp = self._read_resp(timeout = min(remaining, 2.0)) if resp is None: - # Check subprocess health. if not self._ensure_subprocess_alive(): raise RuntimeError("Export subprocess crashed during wait") continue @@ -294,9 +258,7 @@ class ExportOrchestrator: if rtype == "status": message = resp.get("message", "") logger.info("Export subprocess status: %s", message) - # Surface status in the live log panel too so users see - # high-level progress (e.g. "Importing Unsloth...", - # "Loading checkpoint: ...") alongside subprocess output. + # Surface status in the live log panel for high-level progress. if message: self._append_log( { @@ -477,13 +439,10 @@ class ExportOrchestrator: ) def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]: - """Send an export command to the subprocess and wait for result. + """Send an export command and wait for the result. - Returns ``(success, message, output_path)``. ``output_path`` is - the resolved on-disk dir the worker wrote to (None when the - export only pushed to Hub or failed before writing). Surfaced via - the export route's ``details.output_path`` so the dialog's success - screen shows where the model landed. + Returns ``(success, message, output_path)``. ``output_path`` is the on-disk + dir the worker wrote to (None if it only pushed to Hub or failed pre-write). """ with self._lock: if not self._ensure_subprocess_alive(): @@ -517,7 +476,6 @@ class ExportOrchestrator: """Cleanup export-related models from memory.""" with self._lock: if not self._ensure_subprocess_alive(): - # No subprocess — clear local state. self.current_checkpoint = None self.is_vision = False self.is_peft = False diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py index a517d6d7b1..fb2a893014 100644 --- a/studio/backend/core/export/worker.py +++ b/studio/backend/core/export/worker.py @@ -1,16 +1,12 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 -""" -Export subprocess entry point. +"""Export subprocess entry point. -Each export session runs in a persistent subprocess (mp.get_context("spawn")), -giving a clean interpreter with no stale module state — solving the -transformers version-switching problem completely. - -The subprocess stays alive while a model is loaded, accepting commands -(load, export_merged, export_base, export_gguf, export_lora, cleanup, -shutdown) via mp.Queue. +Each export session runs in a persistent subprocess (mp spawn), giving a clean +interpreter with no stale module state, which solves transformers version +switching. The subprocess stays alive while a model is loaded, accepting commands +(load, export_*, cleanup, shutdown) via mp.Queue. Pattern follows core/inference/worker.py and core/training/worker.py. """ @@ -31,40 +27,31 @@ from typing import Any logger = get_logger(__name__) -# Gate controlling whether captured stdout/stderr lines are forwarded to -# the parent's resp_queue (and on to the export-dialog SSE stream). Closed -# by default so the noisy bootstrap phase -- transformers venv activation, -# Unsloth/torch imports, base-model resolution, "Top GGUF/hub models" lists, -# vision detection, weight loading bars -- is suppressed in the UI. -# _handle_export() opens the gate when actual export work starts and leaves -# it open; the orchestrator spawns a fresh subprocess for the next checkpoint -# load (see orchestrator._spawn_subprocess), resetting this state. -# -# Lines dropped while the gate is closed are still echoed to the saved -# original stdout/stderr fds so the server console / log file keeps the full -# output for debugging. +# Gate controlling whether captured stdout/stderr lines are forwarded to the +# parent's resp_queue (and on to the export-dialog SSE stream). Closed by default +# so the noisy bootstrap phase (imports, model resolution, loading bars) is +# suppressed in the UI; _handle_export() opens it when export work starts. The +# orchestrator spawns a fresh subprocess per checkpoint load, resetting this. +# Dropped lines are still echoed to the saved fds so the server log keeps them. _log_forward_gate = threading.Event() def _setup_log_capture(resp_queue: Any) -> None: - """Redirect fds 1 and 2 through pipes so every line printed by this - worker and any child it spawns is forwarded to the parent via resp_queue - as {"type": "log", ...} messages. + """Redirect fds 1 and 2 through pipes so every line printed by this worker + and any child it spawns is forwarded to the parent via resp_queue as + {"type": "log", ...} messages. Must run BEFORE LogConfig.setup_logging and any ML imports, else library handlers may capture the original stderr reference and bypass the pipe. - - Lines are also echoed back to the original stdout/stderr so the server - console keeps the full subprocess output, even while - ``_log_forward_gate`` is closed. + Lines are also echoed back to the original fds so the server console keeps + the full output even while ``_log_forward_gate`` is closed. """ try: saved_out_fd = os.dup(1) saved_err_fd = os.dup(2) except OSError: - # dup failed (exotic platforms) - give up quietly; export still - # works, just no live log streaming. + # dup failed; give up quietly (export still works, no live streaming). return try: @@ -86,13 +73,11 @@ def _setup_log_capture(resp_queue: Any) -> None: pass return - # Close the write ends we just dup2'd (fds 1 and 2 are now the real - # write ends). + # Close the write ends we just dup2'd (fds 1 and 2 are the real write ends). os.close(w_out) os.close(w_err) - # Replace sys.stdout/sys.stderr with line-buffered writers bound to the - # (now-redirected) fds 1 and 2. + # Replace sys.stdout/sys.stderr with line-buffered writers on fds 1 and 2. try: sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace") sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace") @@ -110,8 +95,7 @@ def _setup_log_capture(resp_queue: Any) -> None: continue if not chunk: break - # Echo to the original fd so the server console keeps the - # full output. + # Echo to the original fd so the server console keeps the full output. try: os.write(echo_fd, chunk) except OSError: @@ -131,9 +115,8 @@ def _setup_log_capture(resp_queue: Any) -> None: if not line: continue if not _log_forward_gate.is_set(): - # Gate closed (bootstrap) -- already echoed to the saved - # console fd above; drop the line so the export dialog - # doesn't see import / vendoring noise. + # Gate closed (bootstrap): already echoed above; drop the + # line so the export dialog skips import noise. continue try: resp_queue.put_nowait( @@ -145,8 +128,7 @@ def _setup_log_capture(resp_queue: Any) -> None: } ) except Exception: - # Queue put failed (full, closed, etc.) - drop the line - # rather than crash the reader thread. + # Queue put failed; drop the line rather than crash the thread. pass if buf and _log_forward_gate.is_set(): try: @@ -179,7 +161,7 @@ def _setup_log_capture(resp_queue: Any) -> None: def _activate_transformers_version(model_name: str) -> None: """Activate the correct transformers version BEFORE any ML imports.""" - # Ensure backend is on sys.path for utils imports + # Ensure backend is on sys.path for utils imports. backend_path = str(Path(__file__).resolve().parent.parent.parent) if backend_path not in sys.path: sys.path.insert(0, backend_path) @@ -265,11 +247,9 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None: export_type = cmd["export_type"] # "merged", "base", "gguf", "lora" response_type = f"export_{export_type}_done" - # Open the log forwarding gate so the user sees actual export progress - # (Unsloth merge bars, file copies, GGUF conversion, etc.) in the live - # log panel. The gate stays open for the rest of this subprocess's life; - # the orchestrator spawns a fresh subprocess for the next checkpoint - # load, which resets the gate to closed. + # Open the log forwarding gate so the user sees export progress in the live + # log panel. Stays open for the rest of this subprocess's life; the + # orchestrator spawns a fresh subprocess per checkpoint load, resetting it. _log_forward_gate.set() output_path: Any = None @@ -370,23 +350,16 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None """ import queue as _queue - # Install fd-level stdout/stderr capture FIRST so every subsequent print - # and every child process inherits the redirected fds. This powers the - # live export log stream in the UI. + # Install fd-level stdout/stderr capture FIRST so every subsequent print and + # every child process inherits the redirected fds (powers the live log stream). _setup_log_capture(resp_queue) os.environ["TOKENIZERS_PARALLELISM"] = "false" - os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports - # Force unbuffered output from child Python processes (e.g. the GGUF - # converter) so their prints surface in the log stream as they happen, - # not at the end. + os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports + # Unbuffered output from child Python (e.g. GGUF converter) so prints surface live. os.environ["PYTHONUNBUFFERED"] = "1" - # tqdm defaults to a 10-second mininterval when stdout is not a tty - # (it isn't -- we redirected fd 1/2 to a pipe), making multi-step - # progress bars look frozen in the export log panel. Force frequent - # flushes so the user sees movement during merge / GGUF conversion. No - # effect on single-step bars (e.g. "Copying 1 files") which only emit - # start/end events anyway. + # tqdm defaults to a 10s mininterval when stdout isn't a tty (we redirected + # fd 1/2 to a pipe), making multi-step bars look frozen; force frequent flushes. os.environ.setdefault("TQDM_MININTERVAL", "0.5") import warnings @@ -430,10 +403,8 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None ) # ── 1c. Stub torchao on Windows ROCm ── - # Shared with the training worker; see core/_torchao_stub.py for the full - # rationale (torchao -> torch.distributed._functional_collectives crashes - # on Windows ROCm: RCCL backend absent). No-op off Windows ROCm. Must run - # before importing transformers / unsloth_zoo. + # See core/_torchao_stub.py: torchao crashes on Windows ROCm (RCCL absent). + # No-op off Windows ROCm. Must run before importing transformers / unsloth_zoo. from core._torchao_stub import install_torchao_windows_rocm_stub install_torchao_windows_rocm_stub() @@ -509,7 +480,7 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None try: if cmd_type == "load": - # Load a new checkpoint, reusing this subprocess + # Load a new checkpoint, reusing this subprocess. backend.cleanup_memory() _handle_load(backend, cmd, resp_queue) diff --git a/studio/backend/core/inference/_html_to_md.py b/studio/backend/core/inference/_html_to_md.py index b6262d9974..e7de4a5312 100644 --- a/studio/backend/core/inference/_html_to_md.py +++ b/studio/backend/core/inference/_html_to_md.py @@ -81,9 +81,7 @@ class _MarkdownRenderer(HTMLParser): self._pre_parts: list[str] = [] self._in_inline_code: bool = False - # Blockquote state -- stack of output buffers so nested blockquotes - # each collect their own content and get prefixed with the correct - # number of ">" markers on close. + # Blockquote state: stack of buffers so nested blockquotes get the right ">" depth. self._bq_stack: list[list[str]] = [] # ------------------------------------------------------------------ @@ -116,10 +114,7 @@ class _MarkdownRenderer(HTMLParser): prefixed.append(">") return "\n".join(prefixed) - # ------------------------------------------------------------------ - # Table helpers -- flush open cells/rows so HTML with omitted optional - # end tags (, ) does not lose data. - # ------------------------------------------------------------------ + # Table helpers: flush open cells/rows so omitted / don't lose data. def _finish_cell(self) -> None: if not self._in_cell: return @@ -142,10 +137,7 @@ class _MarkdownRenderer(HTMLParser): self._current_row = [] self._row_has_th = False - # ------------------------------------------------------------------ - # Link text helper -- normalize whitespace so block-level content inside - # an does not produce multiline Markdown link labels. - # ------------------------------------------------------------------ + # Link text helper: normalize whitespace so block content in stays single-line. def _finish_link(self) -> None: text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip() href = self._link_href or "" @@ -234,22 +226,19 @@ class _MarkdownRenderer(HTMLParser): self._emit("\n\n") elif tag == "tr": - # Flush any open cell/row from a previous row that omitted its - # optional or end tags. + # Flush open cell/row from a prior row that omitted /. self._finish_cell() self._finish_row() elif tag in ("th", "td"): - # Flush any open cell (handles omitted /). - self._finish_cell() + self._finish_cell() # handles omitted / self._cell_parts = [] self._in_cell = True if tag == "th": self._row_has_th = True elif tag == "img": - # Skip images -- keeps page text focused on readable content and - # avoids data-URI amplification. + # Skip images: keeps text readable, avoids data-URI amplification. return def handle_endtag(self, tag: str) -> None: @@ -310,7 +299,7 @@ class _MarkdownRenderer(HTMLParser): self._finish_row() elif tag == "table": - # Flush any remaining row (handles omitted ). + # Flush remaining row (handles omitted ). self._finish_cell() self._finish_row() self._in_table = False @@ -331,9 +320,7 @@ class _MarkdownRenderer(HTMLParser): return # Collapse all whitespace (including newlines) per HTML rules. text = re.sub(r"\s+", " ", data) - # Suppress whitespace-only text nodes between table structural elements - # (source indentation) so leading spaces don't break Markdown table - # row alignment. + # Suppress whitespace-only nodes between table elements (source indentation). if self._in_table and not self._in_cell and not text.strip(): return self._emit(text) @@ -348,18 +335,10 @@ class _MarkdownRenderer(HTMLParser): return self._emit(html.unescape(f"&#{name};")) - # ------------------------------------------------------------------ # Flush pending buffers (handles truncated HTML from capped fetches) - # ------------------------------------------------------------------ def flush_pending(self) -> None: - """Flush any open side-buffers into ``_out``. - - Called after ``close()`` to recover content from truncated HTML where - closing tags were never seen (common when ``_fetch_page_text`` caps the - download by byte count). - """ + """Flush open side-buffers into ``_out`` after close(), recovering truncated HTML.""" # Flush innermost buffers first so their content propagates outward. - if self._in_link: self._finish_link() @@ -388,15 +367,9 @@ class _MarkdownRenderer(HTMLParser): self._out.append("\n\n" + prefixed + "\n\n") -# ------------------------------------------------------------------ # Post-processing -# ------------------------------------------------------------------ def _cleanup(text: str) -> str: - """Normalize whitespace and blank lines in the final output. - - Preserves content inside fenced code blocks verbatim so intentional blank - lines in ``
`` content are not collapsed.
-    """
+    """Normalize whitespace and blank lines, preserving fenced code blocks verbatim."""
     lines = text.split("\n")
     out: list[str] = []
     in_fence = False
@@ -411,7 +384,6 @@ def _cleanup(text: str) -> str:
             continue
 
         if in_fence:
-            # Preserve code block content exactly as-is.
             out.append(line)
             continue
 
@@ -427,15 +399,11 @@ def _cleanup(text: str) -> str:
     return "\n".join(out).strip()
 
 
-# ------------------------------------------------------------------
 # Public API
-# ------------------------------------------------------------------
 def html_to_markdown(source_html: str) -> str:
-    """Convert an HTML string to Markdown.
+    """Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
 
-    Handles headings, links, bold/italic, ordered/unordered lists, tables,
-    blockquotes, code blocks, and HTML entities. ``