From 8184429243f0dc7ecb4208f6c42080a5a36e0e2a Mon Sep 17 00:00:00 2001
From: Cursor Agent <cursoragent@cursor.com>
Date: Wed, 23 Sep 2026 08:08:34 +0000
Subject: [PATCH] fix(novita): longest-first ARG expand and readiness cleanup

Prevent unbraced $BASE from corrupting $BASE_IMAGE during Dockerfile
FROM rewriting, and keep wait_for_ready inside try/finally in the
simple TB2 example so a readiness timeout still stops the sandbox.

Post-0.6.0 follow-up; not for the Thursday release cut.

Co-authored-by: benjamin.burtenshaw <benjamin.burtenshaw@huggingface.co>
---
 examples/novita_tbench2_simple.py             |  6 ++---
 .../containers/runtime/novita_provider.py     | 13 +++++++---
 tests/test_core/test_novita_provider.py       | 24 +++++++++++++++++++
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/examples/novita_tbench2_simple.py b/examples/novita_tbench2_simple.py
index f3544591d..9666cc93f 100644
--- a/examples/novita_tbench2_simple.py
+++ b/examples/novita_tbench2_simple.py
@@ -27,11 +27,11 @@ async def main() -> int:
     )
     provider = NovitaSandboxProvider()
     # First start builds a Novita template from the Dockerfile, which can take a
-    # few minutes; later starts reuse the cached template.
+    # few minutes; later starts reuse the cached template. Keep readiness inside
+    # the finally so a timeout still stops the sandbox.
     base_url = provider.start_container(image=image)
-    provider.wait_for_ready(base_url, timeout_s=300)
-
     try:
+        provider.wait_for_ready(base_url, timeout_s=300)
         async with Tbench2Env(base_url=base_url, provider=provider) as env:
             result = await env.reset(task_id=task_id)
             print("Instruction head:")
diff --git a/src/openenv/core/containers/runtime/novita_provider.py b/src/openenv/core/containers/runtime/novita_provider.py
index cdaa45a23..28700d3be 100644
--- a/src/openenv/core/containers/runtime/novita_provider.py
+++ b/src/openenv/core/containers/runtime/novita_provider.py
@@ -125,7 +125,9 @@ def _resolve_from_references(content: str) -> str:
     - The parser stores a ``FROM ${BASE_IMAGE}`` line verbatim, so the template
       would be built from a literal image named ``${BASE_IMAGE}``. Only global
       ARGs (declared before the first ``FROM``) are in scope for a ``FROM``
-      line, which is the form every in-repo Dockerfile uses.
+      line, which is the form every in-repo Dockerfile uses. Substitution is
+      longest-name-first with a word-boundary check on unbraced ``$NAME`` forms
+      so ``$BASE`` cannot corrupt ``$BASE_IMAGE``.
     - ``FROM --platform=linux/amd64 python:3.10-slim`` likewise keeps the flag
       as part of the name. Novita builds for its own platform, so the flag is
       dropped rather than propagated.
@@ -147,9 +149,14 @@ def _resolve_from_references(content: str) -> str:
 
         reference = match.group("rest").strip()
         reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
-        for name, value in arg_defaults.items():
+        # Longest names first so an unbraced `$BASE` cannot corrupt
+        # `$BASE_IMAGE` into `<BASE-value>_IMAGE`. Prefer word-boundary
+        # matching for the unbraced form for the same reason.
+        for name, value in sorted(
+            arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
+        ):
             reference = reference.replace(f"${{{name}}}", value)
-            reference = reference.replace(f"${name}", value)
+            reference = re.sub(rf"\${re.escape(name)}(?!\w)", value, reference)
         out.append(f"FROM {reference}")
     return "\n".join(out)
 
diff --git a/tests/test_core/test_novita_provider.py b/tests/test_core/test_novita_provider.py
index 05318beac..da7d99447 100644
--- a/tests/test_core/test_novita_provider.py
+++ b/tests/test_core/test_novita_provider.py
@@ -871,6 +871,30 @@ def test_arg_after_first_from_not_in_scope(self):
         out = _resolve_from_references("FROM python:3.11\nARG X=alpine\nRUN echo hi\n")
         assert "FROM python:3.11" in out
 
+    def test_resolve_unbraced_arg_does_not_prefix_longer_name(self):
+        from openenv.core.containers.runtime.novita_provider import (
+            _resolve_from_references,
+        )
+
+        # Declaring a shorter ARG before a longer one must not turn
+        # `$BASE_IMAGE` into `<BASE-value>_IMAGE`.
+        out = _resolve_from_references(
+            "ARG BASE=python:3.12\n"
+            "ARG BASE_IMAGE=python:3.11-slim\n"
+            "FROM $BASE_IMAGE\n"
+            "RUN echo hi\n"
+        )
+        assert "FROM python:3.11-slim" in out
+        assert "python:3.12_IMAGE" not in out
+
+        braced = _resolve_from_references(
+            "ARG BASE=python:3.12\n"
+            "ARG BASE_IMAGE=python:3.11-slim\n"
+            "FROM ${BASE_IMAGE}\n"
+        )
+        assert "FROM python:3.11-slim" in braced
+        assert "python:3.12_IMAGE" not in braced
+
     def test_flatten_drops_same_path_copy(self):
         from openenv.core.containers.runtime.novita_provider import _flatten_multistage