diff --git a/packages/modal-infra/src/images/primo_overlay.py b/packages/modal-infra/src/images/primo_overlay.py index f09a7f93d..8e326f661 100644 --- a/packages/modal-infra/src/images/primo_overlay.py +++ b/packages/modal-infra/src/images/primo_overlay.py @@ -4,6 +4,8 @@ localized to the smallest possible hook in base.py. """ +import modal + AWS_CLI_VERSION = "2.34.50" AWS_CLI_SHA256 = "0e6f3d4330a0655e2d08f3791a2ee9503bb55accbac5633b839b8e0b66c0e5b5" @@ -47,6 +49,21 @@ def primo_sandbox_create_kwargs(repo_owner: str | None, repo_name: str | None) - } +async def create_primo_sandbox(repo_owner: str | None, repo_name: str | None, **create_kwargs): + """Create a sandbox the Primo way: our entrypoint wrapper plus per-repo runtime. + + Every `modal.Sandbox.create` call in upstream's manager routes through here, + so the fork's patch at each call site is a single line and never reaches + into upstream's `create_kwargs` dict — the spot upstream keeps adding + fields to, and therefore the spot that keeps conflicting. + + Caller-supplied kwargs win, so explicit `cpu`/`memory` from session settings + still override the Core defaults. + """ + kwargs = {**primo_sandbox_create_kwargs(repo_owner, repo_name), **create_kwargs} + return await modal.Sandbox.create.aio(*PRIMO_SANDBOX_COMMAND, **kwargs) + + def apply_primo_postgres_runtime(image): if not hasattr(image, "apt_install"): return image diff --git a/packages/modal-infra/src/sandbox/manager.py b/packages/modal-infra/src/sandbox/manager.py index b489553f9..7c03004de 100644 --- a/packages/modal-infra/src/sandbox/manager.py +++ b/packages/modal-infra/src/sandbox/manager.py @@ -31,11 +31,7 @@ from ..app import app, llm_secrets from ..images.base import base_image -from ..images.primo_overlay import ( - PRIMO_SANDBOX_COMMAND, - apply_primo_postgres_runtime, - primo_sandbox_create_kwargs, -) +from ..images.primo_overlay import apply_primo_postgres_runtime, create_primo_sandbox from .vcs_env import inject_vcs_env_vars log = get_logger("manager") @@ -428,13 +424,12 @@ async def create_sandbox( "timeout": config.timeout_seconds, "workdir": "/workspace", "env": env_vars, - **primo_sandbox_create_kwargs(config.repo_owner, config.repo_name), **_resource_kwargs(config.settings), } if exposed_ports: create_kwargs["encrypted_ports"] = exposed_ports - sandbox = await modal.Sandbox.create.aio(*PRIMO_SANDBOX_COMMAND, **create_kwargs) + sandbox = await create_primo_sandbox(config.repo_owner, config.repo_name, **create_kwargs) modal_object_id = sandbox.object_id code_server_url, ttyd_url, extra_tunnel_urls = await self._resolve_and_setup_tunnels( @@ -522,15 +517,15 @@ async def create_build_sandbox( inject_vcs_env_vars(env_vars, clone_token or None) - sandbox = await modal.Sandbox.create.aio( - *PRIMO_SANDBOX_COMMAND, + sandbox = await create_primo_sandbox( + repo_owner, + repo_name, image=base_image, app=app, secrets=[], timeout=timeout_seconds, workdir="/workspace", env=env_vars, - **primo_sandbox_create_kwargs(repo_owner, repo_name), ) modal_object_id = sandbox.object_id @@ -745,13 +740,12 @@ async def restore_from_snapshot( "timeout": timeout_seconds, "workdir": "/workspace", "env": env_vars, - **primo_sandbox_create_kwargs(repo_owner, repo_name), **_resource_kwargs(settings), } if exposed_ports: create_kwargs["encrypted_ports"] = exposed_ports - sandbox = await modal.Sandbox.create.aio(*PRIMO_SANDBOX_COMMAND, **create_kwargs) + sandbox = await create_primo_sandbox(repo_owner, repo_name, **create_kwargs) modal_object_id = sandbox.object_id code_server_url, ttyd_url, extra_tunnel_urls = await self._resolve_and_setup_tunnels( diff --git a/packages/modal-infra/tests/test_build_sandbox.py b/packages/modal-infra/tests/test_build_sandbox.py index fc0a4fc4b..759369487 100644 --- a/packages/modal-infra/tests/test_build_sandbox.py +++ b/packages/modal-infra/tests/test_build_sandbox.py @@ -4,7 +4,6 @@ import pytest -from src.images.primo_overlay import PRIMO_SANDBOX_COMMAND from src.sandbox.manager import SandboxManager @@ -45,30 +44,6 @@ async def test_env_vars_include_image_build_mode(monkeypatch): assert env["IMAGE_BUILD_MODE"] == "true" -@pytest.mark.asyncio -async def test_starts_postgres_before_sandbox_runtime(monkeypatch): - captured = {} - monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_sandbox_create(captured)) - - manager = SandboxManager() - await manager.create_build_sandbox(repo_owner="acme", repo_name="my-repo") - - assert captured["args"] == PRIMO_SANDBOX_COMMAND - - -@pytest.mark.asyncio -async def test_core_build_uses_vm_runtime_with_ci_sized_resources(monkeypatch): - captured = {} - monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_sandbox_create(captured)) - - manager = SandboxManager() - await manager.create_build_sandbox(repo_owner="primo-devs", repo_name="core") - - assert captured["kwargs"]["cpu"] == 2.0 - assert captured["kwargs"]["memory"] == 8192 - assert captured["kwargs"]["experimental_options"] == {"vm_runtime": True} - - @pytest.mark.asyncio async def test_env_vars_include_repo_info(monkeypatch): """Should include REPO_OWNER, REPO_NAME, and SANDBOX_ID.""" diff --git a/packages/modal-infra/tests/test_primo_sandbox_creation.py b/packages/modal-infra/tests/test_primo_sandbox_creation.py new file mode 100644 index 000000000..09edbcfce --- /dev/null +++ b/packages/modal-infra/tests/test_primo_sandbox_creation.py @@ -0,0 +1,90 @@ +"""Primo's sandbox-creation guarantees, kept out of upstream's test modules. + +Upstream's `test_build_sandbox.py` and `test_sandbox_resources.py` churn on +nearly every sync, so fork assertions live here instead — a file upstream will +never touch and git will never have to merge. +""" + +from unittest.mock import AsyncMock + +import pytest + +from src.images.primo_overlay import PRIMO_SANDBOX_COMMAND +from src.sandbox.manager import SandboxConfig, SandboxManager + + +def _fake_create(captured: dict): + """Fake `Sandbox.create` that records the argv and kwargs it was called with.""" + + async def fake_create_aio(*args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + + class FakeSandbox: + object_id = "obj-primo-1" + stdout = None + + return FakeSandbox() + + fake_create_aio.aio = fake_create_aio + return fake_create_aio + + +@pytest.mark.asyncio +async def test_build_sandbox_starts_postgres_before_sandbox_runtime(monkeypatch): + captured: dict = {} + monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_create(captured)) + + await SandboxManager().create_build_sandbox(repo_owner="acme", repo_name="my-repo") + + assert captured["args"] == PRIMO_SANDBOX_COMMAND + + +@pytest.mark.asyncio +async def test_build_sandbox_for_core_uses_vm_runtime_with_ci_sized_resources(monkeypatch): + captured: dict = {} + monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_create(captured)) + + await SandboxManager().create_build_sandbox(repo_owner="primo-devs", repo_name="core") + + assert captured["kwargs"]["cpu"] == 2.0 + assert captured["kwargs"]["memory"] == 8192 + assert captured["kwargs"]["experimental_options"] == {"vm_runtime": True} + + +@pytest.mark.asyncio +async def test_session_sandbox_starts_postgres_before_sandbox_runtime(monkeypatch): + captured: dict = {} + monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_create(captured)) + monkeypatch.setattr( + SandboxManager, + "_resolve_and_setup_tunnels", + AsyncMock(return_value=(None, None, None)), + ) + + await SandboxManager().create_sandbox(SandboxConfig(repo_owner="acme", repo_name="my-repo")) + + assert captured["args"] == PRIMO_SANDBOX_COMMAND + + +@pytest.mark.asyncio +async def test_session_sandbox_for_core_lets_explicit_resources_override_vm_defaults(monkeypatch): + captured: dict = {} + monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_create(captured)) + monkeypatch.setattr( + SandboxManager, + "_resolve_and_setup_tunnels", + AsyncMock(return_value=(None, None, None)), + ) + + await SandboxManager().create_sandbox( + SandboxConfig( + repo_owner="primo-devs", + repo_name="core", + settings={"cpuCores": 3, "memoryMib": 6144}, + ) + ) + + assert captured["kwargs"]["cpu"] == 3.0 + assert captured["kwargs"]["memory"] == 6144 + assert captured["kwargs"]["experimental_options"] == {"vm_runtime": True} diff --git a/packages/modal-infra/tests/test_sandbox_resources.py b/packages/modal-infra/tests/test_sandbox_resources.py index 730e0924b..2ac81e90e 100644 --- a/packages/modal-infra/tests/test_sandbox_resources.py +++ b/packages/modal-infra/tests/test_sandbox_resources.py @@ -87,29 +87,6 @@ async def test_create_sandbox_omits_resources_without_settings(self, monkeypatch assert "cpu" not in captured["kwargs"] assert "memory" not in captured["kwargs"] - @pytest.mark.asyncio - async def test_core_uses_vm_defaults_and_explicit_resources_override_them(self, monkeypatch): - captured: dict = {} - monkeypatch.setattr("src.sandbox.manager.modal.Sandbox.create", _fake_create(captured)) - monkeypatch.setattr( - SandboxManager, - "_resolve_and_setup_tunnels", - AsyncMock(return_value=(None, None, None)), - ) - - manager = SandboxManager() - await manager.create_sandbox( - SandboxConfig( - repo_owner="primo-devs", - repo_name="core", - settings={"cpuCores": 3, "memoryMib": 6144}, - ) - ) - - assert captured["kwargs"]["cpu"] == 3.0 - assert captured["kwargs"]["memory"] == 6144 - assert captured["kwargs"]["experimental_options"] == {"vm_runtime": True} - @pytest.mark.asyncio async def test_restore_from_snapshot_passes_resources(self, monkeypatch): captured: dict = {} diff --git a/packages/slack-bot/src/classifier/index.primo.test.ts b/packages/slack-bot/src/classifier/index.primo.test.ts new file mode 100644 index 000000000..1f3d91a2a --- /dev/null +++ b/packages/slack-bot/src/classifier/index.primo.test.ts @@ -0,0 +1,117 @@ +/** + * Primo's classifier-prompt guarantees, kept out of upstream's `index.test.ts`. + * + * That file churns on nearly every sync, so fork assertions live here instead — + * a file upstream will never touch and git will never have to merge. + */ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import type { Env, RepoConfig } from "../types"; + +const { + mockMessagesCreate, + mockGetAvailableRepos, + mockBuildRepoDescriptions, + mockGetRoutingRules, + mockGetAvailableEnvironments, +} = vi.hoisted(() => ({ + mockMessagesCreate: vi.fn(), + mockGetAvailableRepos: vi.fn(), + mockBuildRepoDescriptions: vi.fn(), + mockGetRoutingRules: vi.fn(), + mockGetAvailableEnvironments: vi.fn(), +})); + +vi.mock("@anthropic-ai/sdk", () => ({ + // vitest 4 only treats `function`/`class` implementations as constructable; + // an arrow function here throws "is not a constructor" on `new Anthropic()`. + default: vi.fn().mockImplementation(function () { + return { messages: { create: mockMessagesCreate } }; + }), +})); + +vi.mock("./repos", () => ({ + getAvailableRepos: mockGetAvailableRepos, + buildRepoDescriptions: mockBuildRepoDescriptions, + getRoutingRules: mockGetRoutingRules, +})); + +vi.mock("./environments", async (importOriginal) => ({ + ...((await importOriginal()) as object), + getAvailableEnvironments: mockGetAvailableEnvironments, + getEnvironmentById: vi.fn(), +})); + +import { RepoClassifier } from "./index"; +import { PRIMO_CLASSIFIER_INSTRUCTIONS } from "./primo-classifier-instructions"; + +const TEST_REPOS: RepoConfig[] = [ + { + id: "acme/prod", + owner: "acme", + name: "prod", + fullName: "acme/prod", + displayName: "prod", + description: "Production worker", + defaultBranch: "main", + private: true, + aliases: ["production"], + keywords: ["worker", "slack"], + }, + { + id: "acme/web", + owner: "acme", + name: "web", + fullName: "acme/web", + displayName: "web", + description: "Web application", + defaultBranch: "main", + private: true, + aliases: ["frontend"], + keywords: ["react", "ui"], + }, +]; + +const TEST_ENV = { + ANTHROPIC_API_KEY: "test-api-key", + CLASSIFICATION_MODEL: "claude-haiku-4-5", +} as Env; + +describe("RepoClassifier (Primo)", () => { + beforeEach(() => { + vi.clearAllMocks(); + mockGetAvailableRepos.mockResolvedValue(TEST_REPOS); + mockGetRoutingRules.mockResolvedValue([]); + mockGetAvailableEnvironments.mockResolvedValue([]); + mockBuildRepoDescriptions.mockResolvedValue("- acme/prod\n- acme/web"); + mockMessagesCreate.mockResolvedValue({ + content: [ + { + type: "tool_use", + id: "toolu_primo", + name: "classify_target", + input: { + targetId: "acme/prod", + confidence: "high", + reasoning: "Defaulted to the core repository.", + alternatives: [], + }, + }, + ], + }); + }); + + it("adds the Primo default-repository instructions to the LLM prompt", async () => { + const classifier = new RepoClassifier(TEST_ENV); + await classifier.classify("estas vivo infeliz?", undefined, "trace-primo"); + + expect(mockMessagesCreate).toHaveBeenCalledWith( + expect.objectContaining({ + messages: [ + expect.objectContaining({ + content: expect.stringContaining(PRIMO_CLASSIFIER_INSTRUCTIONS.trim()), + }), + ], + }) + ); + }); +}); diff --git a/packages/slack-bot/src/classifier/index.test.ts b/packages/slack-bot/src/classifier/index.test.ts index 95753a738..651a7d028 100644 --- a/packages/slack-bot/src/classifier/index.test.ts +++ b/packages/slack-bot/src/classifier/index.test.ts @@ -139,37 +139,6 @@ describe("RepoClassifier", () => { ); }); - it("adds Primo default repository instructions to the LLM prompt", async () => { - mockMessagesCreate.mockResolvedValue({ - content: [ - { - type: "tool_use", - id: "toolu_primo", - name: "classify_repository", - input: { - repoId: "acme/prod", - confidence: "high", - reasoning: "Defaulted to core-equivalent repo.", - alternatives: [], - }, - }, - ], - }); - - const classifier = new RepoClassifier(TEST_ENV); - await classifier.classify("estas vivo infeliz?", undefined, "trace-primo"); - - expect(mockMessagesCreate).toHaveBeenCalledWith( - expect.objectContaining({ - messages: [ - expect.objectContaining({ - content: expect.stringContaining('repository named "core"'), - }), - ], - }) - ); - }); - it("asks for clarification when tool payload is invalid", async () => { mockMessagesCreate.mockResolvedValue({ content: [