Harden Jinja2 prompt rendering with ImmutableSandboxedEnvironment (CWE-1336 defense-in-depth)

sdk/eidolon_ai_sdk/agent/retriever_agent/hyde_question_transformer.py

@@ -1,6 +1,7 @@
 from typing import List
 
-from jinja2 import Environment, StrictUndefined
+from jinja2 import StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import Field
 
 from eidolon_ai_sdk.agent.retriever_agent.question_transformer import QuestionTransformerSpec, QuestionTransformer
@@ -22,7 +23,7 @@ def __init__(self, **kwargs):
 
     async def transform(self, apu: APU, process_id: str, question: str) -> List[str]:
         thread = apu.new_thread(process_id)
-        env = Environment(undefined=StrictUndefined)
+        env = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
         userPrompt = env.from_string(self.spec.prompt).render(question=question)
         response = await thread.run_request(prompts=[UserTextAPUMessage(prompt=userPrompt)], output_format="str")
 

sdk/eidolon_ai_sdk/agent/retriever_agent/multi_question_transformer.py

@@ -1,6 +1,7 @@
 from typing import List
 
-from jinja2 import Environment, StrictUndefined
+from jinja2 import StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import Field, BaseModel
 
 from eidolon_ai_sdk.agent.retriever_agent.question_transformer import QuestionTransformerSpec, QuestionTransformer
@@ -31,7 +32,7 @@ def __init__(self, **kwargs):
 
     async def transform(self, apu: APU, process_id: str, question: str) -> List[str]:
         thread = apu.new_thread(process_id)
-        env = Environment(undefined=StrictUndefined)
+        env = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
         userPrompt = env.from_string(self.spec.prompt).render(
             question=question, number_to_generate=self.spec.number_to_generate
         )

sdk/eidolon_ai_sdk/agent/simple_agent.py

@@ -6,6 +6,7 @@
 
 from fastapi import Body
 from jinja2 import Environment, meta, StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import field_validator, model_validator, BaseModel, Field
 from pydantic_core import to_jsonable_python, SchemaError
 
@@ -138,7 +139,7 @@ async def create_process(self, process_id: str):
         t = default_apu.main_thread(process_id)
 
         template_args = {k: v.instantiate().value for k, v in self.prompt_templates.items()}
-        env = Environment(undefined=StrictUndefined)
+        env = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
         system_prompt = env.from_string(self.system_prompt).render(**template_args)
         await t.set_boot_messages(prompts=[SystemAPUMessage(prompt=system_prompt)])
 
@@ -224,7 +225,7 @@ async def action_fn(process_id, action, **kwargs):
             if isinstance(request_body, dict):
                 template_args.update(request_body)
 
-            env = Environment(undefined=StrictUndefined)
+            env = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
             text_message = UserTextAPUMessage(prompt=env.from_string(action.user_prompt).render(**template_args))
 
             if execute_on_apu and execute_on_apu in apus:

sdk/eidolon_ai_sdk/agent/sql_agent/agent.py

@@ -1,6 +1,7 @@
 from typing import Literal, Optional, cast
 
-from jinja2 import Environment, StrictUndefined, Template
+from jinja2 import StrictUndefined, Template
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import BaseModel, Field
 
 from eidolon_ai_client.events import (
@@ -72,7 +73,7 @@ def __init__(self, **kwargs):
         else:
             raise ValueError("APU does not have logic units")
 
-        environment = Environment(undefined=StrictUndefined)
+        environment = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
         self._error_prompt = environment.from_string(self.spec.error_prompt)
         self._system_prompt = environment.from_string(self.spec.system_prompt)
         self._user_prompt = environment.from_string(self.spec.user_prompt)

sdk/eidolon_ai_sdk/agent/tot_agent/checker.py

@@ -1,6 +1,7 @@
 from typing import List
 
-from jinja2 import StrictUndefined, Environment
+from jinja2 import StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import Field, BaseModel
 
 from eidolon_ai_sdk.apu.apu import APU
@@ -43,7 +44,7 @@ async def evaluate(
         """
 
         checker_prompt = (
-            Environment(undefined=StrictUndefined)
+            ImmutableSandboxedEnvironment(undefined=StrictUndefined)
             .from_string(self.spec.prompt)
             .render(
                 problem=problem_description,

sdk/eidolon_ai_sdk/agent/tot_agent/thought_generators.py

@@ -6,10 +6,12 @@
 non-repeating thoughts, which are crucial for problem-solving tasks that require
 exploration.
 """
+
 from abc import abstractmethod
 from typing import Any, Dict, List, Callable, Awaitable
 
-from jinja2 import StrictUndefined, Environment
+from jinja2 import StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import Field, BaseModel
 
 from eidolon_ai_sdk.apu.agent_io import SystemAPUMessage, UserTextAPUMessage, APUMessageTypes
@@ -36,7 +38,7 @@ class ThoughtGenerationStrategy(Specable[TGSConfig]):
     """
 
     spec: TGSConfig
-    env = Environment(undefined=StrictUndefined)
+    env = ImmutableSandboxedEnvironment(undefined=StrictUndefined)
 
     def __init__(self, spec):
         super().__init__(spec)

sdk/eidolon_ai_sdk/agent/tot_agent/tot_agent.py

@@ -4,7 +4,8 @@
 from typing import List, Dict, Any, Literal, Optional, Type, Union
 
 from fastapi import HTTPException
-from jinja2 import StrictUndefined, Environment
+from jinja2 import StrictUndefined
+from jinja2.sandbox import ImmutableSandboxedEnvironment
 from pydantic import Field, BaseModel
 
 from eidolon_ai_sdk.agent.agent import register_program
@@ -121,7 +122,11 @@ async def question(self, process_id, body) -> TotResponse:
         # override to run the tree of thoughts algorithm in a separate thread
         thoughts_path: List[str] = []
         level = 0
-        question = Environment(undefined=StrictUndefined).from_string(self.spec.user_prompt).render(**body.model_dump())
+        question = (
+            ImmutableSandboxedEnvironment(undefined=StrictUndefined)
+            .from_string(self.spec.user_prompt)
+            .render(**body.model_dump())
+        )
 
         async def exec_request(
             _boot_messages: List[LLMMessage],