Update metadata_search/research_synthesis.py

metadata_search/research_synthesis.py

@@ -214,19 +214,41 @@ def _research_generate_answer(
         max_new_tokens,
     )
 
-    # _generate_answer normally calls an imported agent_answer function.
-    # Clone that function too so its os.getenv / token constants are isolated.
-    original_agent_answer = generate_globals.get("agent_answer")
-    if isinstance(original_agent_answer, types.FunctionType):
-        agent_globals = _override_token_globals(
-            original_agent_answer.__globals__,
+# _generate_answer normally calls an imported agent_answer function.
+# Clone agent_answer and its run_agent dependency so the research-only
+# token override reaches the actual /agent_run call.
+original_agent_answer = generate_globals.get("agent_answer")
+
+if isinstance(original_agent_answer, types.FunctionType):
+    agent_globals = _override_token_globals(
+        original_agent_answer.__globals__,
+        max_new_tokens,
+    )
+
+    # agent_answer() calls run_agent(), and run_agent() is where
+    # WATERSCOPE_MAX_TOKENS is actually read.
+    original_run_agent = agent_globals.get("run_agent")
+
+    if isinstance(original_run_agent, types.FunctionType):
+        run_agent_globals = _override_token_globals(
+            original_run_agent.__globals__,
             max_new_tokens,
         )
-        cloned_agent_answer = _copy_function(
-            original_agent_answer,
-            agent_globals,
+
+        cloned_run_agent = _copy_function(
+            original_run_agent,
+            run_agent_globals,
         )
-        generate_globals["agent_answer"] = cloned_agent_answer
+
+        agent_globals["run_agent"] = cloned_run_agent
+
+    cloned_agent_answer = _copy_function(
+        original_agent_answer,
+        agent_globals,
+    )
+
+    generate_globals["agent_answer"] = cloned_agent_answer
+    
 
     cloned_generate = _copy_function(
         original_generate,