add fallback path when LLM response missing base image

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,14 +75,28 @@ 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(
                     content="""Please wrap the image name in a block like <image>ubuntu:20.04</image> to indicate your choice."""
                 )
             )
 
-    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,

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: