From 8f5b5ff5c4b380ab071ff42221682efd732220c5 Mon Sep 17 00:00:00 2001
From: Daniel Zayas <danielzayas@Daniels-MacBook-Pro.local>
Date: Thu, 4 Dec 2025 11:33:27 -0800
Subject: [PATCH 1/2] update base_image.py to add fallback when LLM response
 missing image

---
 launch/launch/agent/base_image.py | 32 +++++++++++++++++++++++++++++--
 1 file changed, 30 insertions(+), 2 deletions(-)

diff --git a/launch/launch/agent/base_image.py b/launch/launch/agent/base_image.py
index 994c0fe..bbed401 100644
--- a/launch/launch/agent/base_image.py
+++ b/launch/launch/agent/base_image.py
@@ -7,6 +7,18 @@
 from launch.utilities.language_handlers import get_language_handler
 
 
+def _version_key(image: str) -> tuple[int, ...]:
+    """
+    Provide a tuple that increases with the semantic version embedded in the tag.
+    Falls back to (0,) if no numeric components are found.
+    """
+    _, _, version_part = image.partition(":")
+    if not version_part:
+        version_part = image
+    tokens = [int(token) for token in re.findall(r"\d+", version_part)]
+    return tuple(tokens) if tokens else (0,)
+
+
 @auto_catch
 def select_base_image(state: AgentState) -> dict:
     """
@@ -46,7 +58,9 @@ def select_base_image(state: AgentState) -> dict:
     ]
     base_image = None
     trials = 0
-    while not base_image or trials < 5:
+    max_trials = 5
+    last_response_text = ""
+    while trials < max_trials and not base_image:
         trials += 1
         response = llm.invoke(messages)
         if "<image>" in response.content:
@@ -61,6 +75,10 @@ def select_base_image(state: AgentState) -> dict:
                 )
             )
         else:
+            logger.info(
+                "Base image response missing <image> tag, retrying. Raw response: %s",
+                response_text[:500],
+            )
             messages.append(response)
             messages.append(
                 HumanMessage(
@@ -68,7 +86,17 @@ def select_base_image(state: AgentState) -> dict:
                 )
             )
 
-    logger.info(f"Selected base image: {base_image}")
+    if not base_image:
+        fallback = max(candidate_images, key=_version_key) if candidate_images else None
+        base_image = fallback
+        logger.warning(
+            "Base image selection failed after %s trials, defaulting to %s. Last response: %s",
+            max_trials,
+            base_image,
+            last_response_text[:500],
+        )
+    else:
+        logger.info(f"Selected base image: {base_image}")
     return {
         "messages": messages,
         "base_image": base_image,

From 4e607678b09c455056dae8e294512c60752a6a83 Mon Sep 17 00:00:00 2001
From: Daniel Zayas <danielzayas@Daniels-MacBook-Pro.local>
Date: Thu, 4 Dec 2025 12:20:04 -0800
Subject: [PATCH 2/2] update setup.py to add fallback when LLM response missing
 image

---
 launch/launch/agent/setup.py | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/launch/launch/agent/setup.py b/launch/launch/agent/setup.py
index 1bf26f4..3cf06bd 100644
--- a/launch/launch/agent/setup.py
+++ b/launch/launch/agent/setup.py
@@ -127,9 +127,22 @@ def start_bash_session(state: AgentState) -> dict:
     Returns:
         dict: Updated state with session and pypiserver
     """
+    logger = state["logger"]
+    language = state["language"]
+    language_handler = get_language_handler(language)
+
     base_image = state["base_image"]
+    if not base_image:
+        fallback_image = (
+            language_handler.base_images[-1] if language_handler.base_images else None
+        )
+        base_image = fallback_image
+        logger.warning(
+            "Base image missing from state, defaulting to %s for language %s",
+            base_image,
+            language,
+        )
     repo_root = state["repo_root"]
-    logger = state["logger"]
     logger.info(f"Starting bash session in container based on image: {base_image}")
     session = start_session(base_image, state["instance"])
     logger.info(f"Session started: {session}")
@@ -139,9 +152,6 @@ def start_bash_session(state: AgentState) -> dict:
     logger.info(f"Repo root in the host cleaned up: {repo_root}")
 
     # Setup language-specific environment
-    language = state["language"]
-    language_handler = get_language_handler(language)
-    
     logger.info(f"Setting up environment for language: {language}")
     server = language_handler.setup_environment(session, state["date"])
     if server: