letta-nightly 0.8.5.dev20250624104340__py3-none-any.whl → 0.8.6.dev20250625222533__py3-none-any.whl

letta/agent.py

@@ -255,7 +255,7 @@ class Agent(BaseAgent):
         self.tool_rules_solver.register_tool_call(function_name)
 
         # Extend conversation with function response
-        function_response = package_function_response(False, error_msg)
+        function_response = package_function_response(False, error_msg, self.agent_state.timezone)
         new_message = Message(
             agent_id=self.agent_state.id,
             # Base info OpenAI-style
@@ -640,7 +640,7 @@ class Agent(BaseAgent):
                     function_response, return_char_limit=return_char_limit, truncate=truncate
                 )
                 function_args.pop("self", None)
-                function_response = package_function_response(True, function_response_string)
+                function_response = package_function_response(True, function_response_string, self.agent_state.timezone)
                 function_failed = False
             except Exception as e:
                 function_args.pop("self", None)
@@ -763,7 +763,7 @@ class Agent(BaseAgent):
         self.tool_rules_solver.clear_tool_history()
 
         # Convert MessageCreate objects to Message objects
-        next_input_messages = convert_message_creates_to_messages(input_messages, self.agent_state.id)
+        next_input_messages = convert_message_creates_to_messages(input_messages, self.agent_state.id, self.agent_state.timezone)
         counter = 0
         total_usage = UsageStatistics()
         step_count = 0
@@ -823,7 +823,7 @@ class Agent(BaseAgent):
                         model=self.model,
                         openai_message_dict={
                             "role": "user",  # TODO: change to system?
-                            "content": get_heartbeat(FUNC_FAILED_HEARTBEAT_MESSAGE),
+                            "content": get_heartbeat(self.agent_state.timezone, FUNC_FAILED_HEARTBEAT_MESSAGE),
                         },
                     )
                 ]
@@ -836,7 +836,7 @@ class Agent(BaseAgent):
                         model=self.model,
                         openai_message_dict={
                             "role": "user",  # TODO: change to system?
-                            "content": get_heartbeat(REQ_HEARTBEAT_MESSAGE),
+                            "content": get_heartbeat(self.agent_state.timezone, REQ_HEARTBEAT_MESSAGE),
                         },
                     )
                 ]
@@ -1000,11 +1000,12 @@ class Agent(BaseAgent):
             )
             if job_id:
                 for message in all_new_messages:
-                    self.job_manager.add_message_to_job(
-                        job_id=job_id,
-                        message_id=message.id,
-                        actor=self.user,
-                    )
+                    if message.role != "user":
+                        self.job_manager.add_message_to_job(
+                            job_id=job_id,
+                            message_id=message.id,
+                            actor=self.user,
+                        )
 
             return AgentStepResponse(
                 messages=all_new_messages,
@@ -1079,7 +1080,7 @@ class Agent(BaseAgent):
         assert user_message_str and isinstance(
             user_message_str, str
         ), f"user_message_str should be a non-empty string, got {type(user_message_str)}"
-        user_message_json_str = package_user_message(user_message_str)
+        user_message_json_str = package_user_message(user_message_str, self.agent_state.timezone)
 
         # Validate JSON via save/load
         user_message = validate_json(user_message_json_str)
@@ -1142,7 +1143,9 @@ class Agent(BaseAgent):
         remaining_message_count = 1 + len(in_context_messages) - cutoff  # System + remaining
         hidden_message_count = all_time_message_count - remaining_message_count
         summary_message_count = len(message_sequence_to_summarize)
-        summary_message = package_summarize_message(summary, summary_message_count, hidden_message_count, all_time_message_count)
+        summary_message = package_summarize_message(
+            summary, summary_message_count, hidden_message_count, all_time_message_count, self.agent_state.timezone
+        )
         logger.info(f"Packaged into message: {summary_message}")
 
         prior_len = len(in_context_messages_openai)
@@ -1243,6 +1246,7 @@ class Agent(BaseAgent):
         message_manager_size = self.message_manager.size(actor=self.user, agent_id=self.agent_state.id)
         external_memory_summary = compile_memory_metadata_block(
             memory_edit_timestamp=get_utc_time(),
+            timezone=self.agent_state.timezone,
             previous_message_count=self.message_manager.size(actor=self.user, agent_id=self.agent_state.id),
             archival_memory_size=self.agent_manager.passage_size(actor=self.user, agent_id=self.agent_state.id),
         )

letta/agents/base_agent.py

@@ -50,7 +50,9 @@ class BaseAgent(ABC):
         self.logger = get_logger(agent_id)
 
     @abstractmethod
-    async def step(self, input_messages: List[MessageCreate], max_steps: int = DEFAULT_MAX_STEPS) -> LettaResponse:
+    async def step(
+        self, input_messages: List[MessageCreate], max_steps: int = DEFAULT_MAX_STEPS, run_id: Optional[str] = None
+    ) -> LettaResponse:
         """
         Main execution loop for the agent.
         """
@@ -118,6 +120,7 @@ class BaseAgent(ABC):
                 system_prompt=agent_state.system,
                 in_context_memory=agent_state.memory,
                 in_context_memory_last_edit=memory_edit_timestamp,
+                timezone=agent_state.timezone,
                 previous_message_count=num_messages - len(in_context_messages),
                 archival_memory_size=num_archival_memories,
                 tool_rules_solver=tool_rules_solver,

letta/agents/helpers.py

@@ -1,12 +1,15 @@
+import json
 import uuid
 import xml.etree.ElementTree as ET
 from typing import List, Optional, Tuple
 
+from letta.helpers import ToolRulesSolver
 from letta.schemas.agent import AgentState
 from letta.schemas.letta_message import MessageType
 from letta.schemas.letta_response import LettaResponse
 from letta.schemas.letta_stop_reason import LettaStopReason, StopReasonType
 from letta.schemas.message import Message, MessageCreate
+from letta.schemas.tool_execution_result import ToolExecutionResult
 from letta.schemas.usage import LettaUsageStatistics
 from letta.schemas.user import User
 from letta.server.rest_api.utils import create_input_messages
@@ -69,7 +72,8 @@ def _prepare_in_context_messages(
 
     # Create a new user message from the input and store it
     new_in_context_messages = message_manager.create_many_messages(
-        create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor), actor=actor
+        create_input_messages(input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor),
+        actor=actor,
     )
 
     return current_in_context_messages, new_in_context_messages
@@ -106,7 +110,8 @@ async def _prepare_in_context_messages_async(
 
     # Create a new user message from the input and store it
     new_in_context_messages = await message_manager.create_many_messages_async(
-        create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor), actor=actor
+        create_input_messages(input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor),
+        actor=actor,
     )
 
     return current_in_context_messages, new_in_context_messages
@@ -141,7 +146,9 @@ async def _prepare_in_context_messages_no_persist_async(
         current_in_context_messages = await message_manager.get_messages_by_ids_async(message_ids=agent_state.message_ids, actor=actor)
 
     # Create a new user message from the input but dont store it yet
-    new_in_context_messages = create_input_messages(input_messages=input_messages, agent_id=agent_state.id, actor=actor)
+    new_in_context_messages = create_input_messages(
+        input_messages=input_messages, agent_id=agent_state.id, timezone=agent_state.timezone, actor=actor
+    )
 
     return current_in_context_messages, new_in_context_messages
 
@@ -201,3 +208,28 @@ def deserialize_message_history(xml_str: str) -> Tuple[List[str], str]:
 
 def generate_step_id():
     return f"step-{uuid.uuid4()}"
+
+
+def _safe_load_dict(raw: str) -> dict:
+    """Lenient JSON → dict with fallback to eval on assertion failure."""
+    if "}{" in raw:  # strip accidental parallel calls
+        raw = raw.split("}{", 1)[0] + "}"
+    try:
+        data = json.loads(raw)
+        if not isinstance(data, dict):
+            raise AssertionError
+        return data
+    except (json.JSONDecodeError, AssertionError):
+        return json.loads(raw) if raw else {}
+
+
+def _pop_heartbeat(tool_args: dict) -> bool:
+    hb = tool_args.pop("request_heartbeat", False)
+    return str(hb).lower() == "true" if isinstance(hb, str) else bool(hb)
+
+
+def _build_rule_violation_result(tool_name: str, valid: list[str], solver: ToolRulesSolver) -> ToolExecutionResult:
+    hint_lines = solver.guess_rule_violation(tool_name)
+    hint_txt = ("\n** Hint: Possible rules that were violated:\n" + "\n".join(f"\t- {h}" for h in hint_lines)) if hint_lines else ""
+    msg = f"[ToolConstraintError] Cannot call {tool_name}, " f"valid tools include: {valid}.{hint_txt}"
+    return ToolExecutionResult(status="error", func_return=msg)

letta/agents/letta_agent.py

@@ -10,8 +10,15 @@ from opentelemetry.trace import Span
 
 from letta.agents.base_agent import BaseAgent
 from letta.agents.ephemeral_summary_agent import EphemeralSummaryAgent
-from letta.agents.helpers import _create_letta_response, _prepare_in_context_messages_no_persist_async, generate_step_id
-from letta.constants import DEFAULT_MAX_STEPS
+from letta.agents.helpers import (
+    _build_rule_violation_result,
+    _create_letta_response,
+    _pop_heartbeat,
+    _prepare_in_context_messages_no_persist_async,
+    _safe_load_dict,
+    generate_step_id,
+)
+from letta.constants import DEFAULT_MAX_STEPS, NON_USER_MSG_PREFIX
 from letta.errors import ContextWindowExceededError
 from letta.helpers import ToolRulesSolver
 from letta.helpers.datetime_helpers import AsyncTimer, get_utc_time, get_utc_timestamp_ns, ns_to_ms
@@ -27,7 +34,7 @@ from letta.otel.context import get_ctx_attributes
 from letta.otel.metric_registry import MetricRegistry
 from letta.otel.tracing import log_event, trace_method, tracer
 from letta.schemas.agent import AgentState, UpdateAgent
-from letta.schemas.enums import MessageRole
+from letta.schemas.enums import MessageRole, ProviderType
 from letta.schemas.letta_message import MessageType
 from letta.schemas.letta_message_content import OmittedReasoningContent, ReasoningContent, RedactedReasoningContent, TextContent
 from letta.schemas.letta_response import LettaResponse
@@ -43,6 +50,7 @@ from letta.server.rest_api.utils import create_letta_messages_from_llm_response
 from letta.services.agent_manager import AgentManager
 from letta.services.block_manager import BlockManager
 from letta.services.helpers.tool_parser_helper import runtime_override_tool_json_schema
+from letta.services.job_manager import JobManager
 from letta.services.message_manager import MessageManager
 from letta.services.passage_manager import PassageManager
 from letta.services.step_manager import NoopStepManager, StepManager
@@ -55,8 +63,6 @@ from letta.system import package_function_response
 from letta.types import JsonDict
 from letta.utils import log_telemetry, validate_function_response
 
-logger = get_logger(__name__)
-
 
 class LettaAgent(BaseAgent):
 
@@ -66,6 +72,7 @@ class LettaAgent(BaseAgent):
         message_manager: MessageManager,
         agent_manager: AgentManager,
         block_manager: BlockManager,
+        job_manager: JobManager,
         passage_manager: PassageManager,
         actor: User,
         step_manager: StepManager = NoopStepManager(),
@@ -81,6 +88,7 @@ class LettaAgent(BaseAgent):
         # TODO: Make this more general, factorable
         # Summarizer settings
         self.block_manager = block_manager
+        self.job_manager = job_manager
         self.passage_manager = passage_manager
         self.step_manager = step_manager
         self.telemetry_manager = telemetry_manager
@@ -95,6 +103,7 @@ class LettaAgent(BaseAgent):
         self.summarization_agent = None
         self.summary_block_label = summary_block_label
         self.max_summarization_retries = max_summarization_retries
+        self.logger = get_logger(agent_id)
 
         # TODO: Expand to more
         if enable_summarization and model_settings.openai_api_key:
@@ -120,6 +129,7 @@ class LettaAgent(BaseAgent):
         self,
         input_messages: List[MessageCreate],
         max_steps: int = DEFAULT_MAX_STEPS,
+        run_id: Optional[str] = None,
         use_assistant_message: bool = True,
         request_start_timestamp_ns: Optional[int] = None,
         include_return_message_types: Optional[List[MessageType]] = None,
@@ -127,10 +137,11 @@ class LettaAgent(BaseAgent):
         agent_state = await self.agent_manager.get_agent_by_id_async(
             agent_id=self.agent_id, include_relationships=["tools", "memory", "tool_exec_environment_variables"], actor=self.actor
         )
-        _, new_in_context_messages, usage, stop_reason = await self._step(
+        _, new_in_context_messages, stop_reason, usage = await self._step(
             agent_state=agent_state,
             input_messages=input_messages,
             max_steps=max_steps,
+            run_id=run_id,
             request_start_timestamp_ns=request_start_timestamp_ns,
         )
         return _create_letta_response(
@@ -193,7 +204,6 @@ class LettaAgent(BaseAgent):
             response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
 
             # update usage
-            # TODO: add run_id
             usage.step_count += 1
             usage.completion_tokens += response.usage.completion_tokens
             usage.prompt_tokens += response.usage.prompt_tokens
@@ -219,7 +229,7 @@ class LettaAgent(BaseAgent):
             elif response.choices[0].message.content:
                 reasoning = [TextContent(text=response.choices[0].message.content)]  # reasoning placed into content for legacy reasons
             else:
-                logger.info("No reasoning content found.")
+                self.logger.info("No reasoning content found.")
                 reasoning = None
 
             persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
@@ -233,8 +243,11 @@ class LettaAgent(BaseAgent):
                 agent_step_span=agent_step_span,
                 is_final_step=(i == max_steps - 1),
             )
-            self.response_messages.extend(persisted_messages)
-            new_in_context_messages.extend(persisted_messages)
+
+            # TODO (cliandy): handle message contexts with larger refactor and dedupe logic
+            new_message_idx = len(initial_messages) if initial_messages else 0
+            self.response_messages.extend(persisted_messages[new_message_idx:])
+            new_in_context_messages.extend(persisted_messages[new_message_idx:])
             initial_messages = None
             log_event("agent.stream_no_tokens.llm_response.processed")  # [4^]
 
@@ -266,7 +279,7 @@ class LettaAgent(BaseAgent):
                 if include_return_message_types is None or message.message_type in include_return_message_types:
                     yield f"data: {message.model_dump_json()}\n\n"
 
-            MetricRegistry().step_execution_time_ms_histogram.record(step_start - get_utc_timestamp_ns(), get_ctx_attributes())
+            MetricRegistry().step_execution_time_ms_histogram.record(get_utc_timestamp_ns() - step_start, get_ctx_attributes())
 
             if not should_continue:
                 break
@@ -302,6 +315,7 @@ class LettaAgent(BaseAgent):
         agent_state: AgentState,
         input_messages: List[MessageCreate],
         max_steps: int = DEFAULT_MAX_STEPS,
+        run_id: Optional[str] = None,
         request_start_timestamp_ns: Optional[int] = None,
     ) -> Tuple[List[Message], List[Message], Optional[LettaStopReason], LettaUsageStatistics]:
         """
@@ -345,11 +359,11 @@ class LettaAgent(BaseAgent):
 
             response = llm_client.convert_response_to_chat_completion(response_data, in_context_messages, agent_state.llm_config)
 
-            # TODO: add run_id
             usage.step_count += 1
             usage.completion_tokens += response.usage.completion_tokens
             usage.prompt_tokens += response.usage.prompt_tokens
             usage.total_tokens += response.usage.total_tokens
+            usage.run_ids = [run_id] if run_id else None
             MetricRegistry().message_output_tokens.record(
                 response.usage.completion_tokens, dict(get_ctx_attributes(), **{"model.name": agent_state.llm_config.model})
             )
@@ -371,7 +385,7 @@ class LettaAgent(BaseAgent):
             elif response.choices[0].message.omitted_reasoning_content:
                 reasoning = [OmittedReasoningContent()]
             else:
-                logger.info("No reasoning content found.")
+                self.logger.info("No reasoning content found.")
                 reasoning = None
 
             persisted_messages, should_continue, stop_reason = await self._handle_ai_response(
@@ -385,9 +399,12 @@ class LettaAgent(BaseAgent):
                 initial_messages=initial_messages,
                 agent_step_span=agent_step_span,
                 is_final_step=(i == max_steps - 1),
+                run_id=run_id,
             )
-            self.response_messages.extend(persisted_messages)
-            new_in_context_messages.extend(persisted_messages)
+            new_message_idx = len(initial_messages) if initial_messages else 0
+            self.response_messages.extend(persisted_messages[new_message_idx:])
+            new_in_context_messages.extend(persisted_messages[new_message_idx:])
+
             initial_messages = None
             log_event("agent.step.llm_response.processed")  # [4^]
 
@@ -435,7 +452,7 @@ class LettaAgent(BaseAgent):
                 force=False,
             )
 
-        return current_in_context_messages, new_in_context_messages, usage, stop_reason
+        return current_in_context_messages, new_in_context_messages, stop_reason, usage
 
     async def _update_agent_last_run_metrics(self, completion_time: datetime, duration_ms: float) -> None:
         try:
@@ -445,7 +462,7 @@ class LettaAgent(BaseAgent):
                 actor=self.actor,
             )
         except Exception as e:
-            logger.error(f"Failed to update agent's last run metrics: {e}")
+            self.logger.error(f"Failed to update agent's last run metrics: {e}")
 
     @trace_method
     async def step_stream(
@@ -512,12 +529,12 @@ class LettaAgent(BaseAgent):
 
             # TODO: THIS IS INCREDIBLY UGLY
             # TODO: THERE ARE MULTIPLE COPIES OF THE LLM_CONFIG EVERYWHERE THAT ARE GETTING MANIPULATED
-            if agent_state.llm_config.model_endpoint_type == "anthropic":
+            if agent_state.llm_config.model_endpoint_type in [ProviderType.anthropic, ProviderType.bedrock]:
                 interface = AnthropicStreamingInterface(
                     use_assistant_message=use_assistant_message,
                     put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
                 )
-            elif agent_state.llm_config.model_endpoint_type == "openai":
+            elif agent_state.llm_config.model_endpoint_type == ProviderType.openai:
                 interface = OpenAIStreamingInterface(
                     use_assistant_message=use_assistant_message,
                     put_inner_thoughts_in_kwarg=agent_state.llm_config.put_inner_thoughts_in_kwargs,
@@ -590,8 +607,10 @@ class LettaAgent(BaseAgent):
                 agent_step_span=agent_step_span,
                 is_final_step=(i == max_steps - 1),
             )
-            self.response_messages.extend(persisted_messages)
-            new_in_context_messages.extend(persisted_messages)
+            new_message_idx = len(initial_messages) if initial_messages else 0
+            self.response_messages.extend(persisted_messages[new_message_idx:])
+            new_in_context_messages.extend(persisted_messages[new_message_idx:])
+
             initial_messages = None
 
             # log total step time
@@ -916,47 +935,18 @@ class LettaAgent(BaseAgent):
         initial_messages: Optional[List[Message]] = None,
         agent_step_span: Optional["Span"] = None,
         is_final_step: Optional[bool] = None,
+        run_id: Optional[str] = None,
     ) -> Tuple[List[Message], bool, Optional[LettaStopReason]]:
         """
-        Now that streaming is done, handle the final AI response.
-        This might yield additional SSE tokens if we do stalling.
-        At the end, set self._continue_execution accordingly.
+        Handle the final AI response once streaming completes, execute / validate the
+        tool call, decide whether we should keep stepping, and persist state.
         """
-        stop_reason = None
-        # Check if the called tool is allowed by tool name:
-        tool_call_name = tool_call.function.name
-        tool_call_args_str = tool_call.function.arguments
-
-        # Temp hack to gracefully handle parallel tool calling attempt, only take first one
-        if "}{" in tool_call_args_str:
-            tool_call_args_str = tool_call_args_str.split("}{", 1)[0] + "}"
-
-        try:
-            tool_args = json.loads(tool_call_args_str)
-            assert isinstance(tool_args, dict), "tool_args must be a dict"
-        except json.JSONDecodeError:
-            tool_args = {}
-        except AssertionError:
-            tool_args = json.loads(tool_args)
-
-        # Get request heartbeats and coerce to bool
-        request_heartbeat = tool_args.pop("request_heartbeat", False)
-        if is_final_step:
-            stop_reason = LettaStopReason(stop_reason=StopReasonType.max_steps.value)
-            logger.info("Agent has reached max steps.")
-            request_heartbeat = False
-        else:
-            # Pre-emptively pop out inner_thoughts
-            tool_args.pop(INNER_THOUGHTS_KWARG, "")
-
-            # So this is necessary, because sometimes non-structured outputs makes mistakes
-            if not isinstance(request_heartbeat, bool):
-                if isinstance(request_heartbeat, str):
-                    request_heartbeat = request_heartbeat.lower() == "true"
-                else:
-                    request_heartbeat = bool(request_heartbeat)
-
-        tool_call_id = tool_call.id or f"call_{uuid.uuid4().hex[:8]}"
+        # 1.  Parse and validate the tool-call envelope
+        tool_call_name: str = tool_call.function.name
+        tool_call_id: str = tool_call.id or f"call_{uuid.uuid4().hex[:8]}"
+        tool_args = _safe_load_dict(tool_call.function.arguments)
+        request_heartbeat: bool = _pop_heartbeat(tool_args)
+        tool_args.pop(INNER_THOUGHTS_KWARG, None)
 
         log_telemetry(
             self.logger,
@@ -966,16 +956,11 @@ class LettaAgent(BaseAgent):
             tool_call_id=tool_call_id,
             request_heartbeat=request_heartbeat,
         )
-        # Check if tool rule is violated - if so, we'll force continuation
-        tool_rule_violated = tool_call_name not in valid_tool_names
 
+        # 2.  Execute the tool (or synthesize an error result if disallowed)
+        tool_rule_violated = tool_call_name not in valid_tool_names
         if tool_rule_violated:
-            base_error_message = f"[ToolConstraintError] Cannot call {tool_call_name}, valid tools to call include: {valid_tool_names}."
-            violated_rule_messages = tool_rules_solver.guess_rule_violation(tool_call_name)
-            if violated_rule_messages:
-                bullet_points = "\n".join(f"\t- {msg}" for msg in violated_rule_messages)
-                base_error_message += f"\n** Hint: Possible rules that were violated:\n{bullet_points}"
-            tool_execution_result = ToolExecutionResult(status="error", func_return=base_error_message)
+            tool_execution_result = _build_rule_violation_result(tool_call_name, valid_tool_names, tool_rules_solver)
         else:
             tool_execution_result = await self._execute_tool(
                 tool_name=tool_call_name,
@@ -984,51 +969,38 @@ class LettaAgent(BaseAgent):
                 agent_step_span=agent_step_span,
                 step_id=step_id,
             )
+
         log_telemetry(
             self.logger, "_handle_ai_response execute tool finish", tool_execution_result=tool_execution_result, tool_call_id=tool_call_id
         )
 
-        if tool_call_name in ["conversation_search", "conversation_search_date", "archival_memory_search"]:
-            # with certain functions we rely on the paging mechanism to handle overflow
-            truncate = False
-        else:
-            # but by default, we add a truncation safeguard to prevent bad functions from
-            # overflow the agent context window
-            truncate = True
-
-        # get the function response limit
-        target_tool = next((x for x in agent_state.tools if x.name == tool_call_name), None)
-        return_char_limit = target_tool.return_char_limit if target_tool else None
+        # 3.  Prepare the function-response payload
+        truncate = tool_call_name not in {"conversation_search", "conversation_search_date", "archival_memory_search"}
+        return_char_limit = next(
+            (t.return_char_limit for t in agent_state.tools if t.name == tool_call_name),
+            None,
+        )
         function_response_string = validate_function_response(
-            tool_execution_result.func_return, return_char_limit=return_char_limit, truncate=truncate
+            tool_execution_result.func_return,
+            return_char_limit=return_char_limit,
+            truncate=truncate,
         )
-        function_response = package_function_response(
+        self.last_function_response = package_function_response(
             was_success=tool_execution_result.success_flag,
             response_string=function_response_string,
+            timezone=agent_state.timezone,
         )
 
-        # 4. Register tool call with tool rule solver
-        # Resolve whether or not to continue stepping
-        continue_stepping = request_heartbeat
-
-        # Force continuation if tool rule was violated to give the model another chance
-        if tool_rule_violated:
-            continue_stepping = True
-        else:
-            tool_rules_solver.register_tool_call(tool_name=tool_call_name)
-            if tool_rules_solver.is_terminal_tool(tool_name=tool_call_name):
-                if continue_stepping:
-                    stop_reason = LettaStopReason(stop_reason=StopReasonType.tool_rule.value)
-                continue_stepping = False
-            elif tool_rules_solver.has_children_tools(tool_name=tool_call_name):
-                continue_stepping = True
-            elif tool_rules_solver.is_continue_tool(tool_name=tool_call_name):
-                continue_stepping = True
+        # 4.  Decide whether to keep stepping  (<<< focal section simplified)
+        continue_stepping, heartbeat_reason, stop_reason = self._decide_continuation(
+            request_heartbeat=request_heartbeat,
+            tool_call_name=tool_call_name,
+            tool_rule_violated=tool_rule_violated,
+            tool_rules_solver=tool_rules_solver,
+            is_final_step=is_final_step,
+        )
 
-        # 5a. Persist Steps to DB
-        # Following agent loop to persist this before messages
-        # TODO (cliandy): determine what should match old loop w/provider_id, job_id
-        # TODO (cliandy): UsageStatistics and LettaUsageStatistics are used in many places, but are not the same.
+        # 5.  Persist step + messages and propagate to jobs
         logged_step = await self.step_manager.log_step_async(
             actor=self.actor,
             agent_id=agent_state.id,
@@ -1039,11 +1011,10 @@ class LettaAgent(BaseAgent):
             context_window_limit=agent_state.llm_config.context_window,
             usage=usage,
             provider_id=None,
-            job_id=None,
+            job_id=run_id,
             step_id=step_id,
         )
 
-        # 5b. Persist Messages to DB
         tool_call_messages = create_letta_messages_from_llm_response(
             agent_id=agent_state.id,
             model=agent_state.llm_config.model,
@@ -1053,20 +1024,74 @@ class LettaAgent(BaseAgent):
             tool_call_id=tool_call_id,
             function_call_success=tool_execution_result.success_flag,
             function_response=function_response_string,
+            timezone=agent_state.timezone,
             actor=self.actor,
-            add_heartbeat_request_system_message=continue_stepping,
+            continue_stepping=continue_stepping,
+            heartbeat_reason=heartbeat_reason,
             reasoning_content=reasoning_content,
             pre_computed_assistant_message_id=pre_computed_assistant_message_id,
-            step_id=logged_step.id if logged_step else None,  # TODO (cliandy): eventually move over other agent loops
+            step_id=logged_step.id if logged_step else None,
         )
 
         persisted_messages = await self.message_manager.create_many_messages_async(
             (initial_messages or []) + tool_call_messages, actor=self.actor
         )
-        self.last_function_response = function_response
+
+        if run_id:
+            await self.job_manager.add_messages_to_job_async(
+                job_id=run_id,
+                message_ids=[m.id for m in persisted_messages if m.role != "user"],
+                actor=self.actor,
+            )
 
         return persisted_messages, continue_stepping, stop_reason
 
+    def _decide_continuation(
+        self,
+        request_heartbeat: bool,
+        tool_call_name: str,
+        tool_rule_violated: bool,
+        tool_rules_solver: ToolRulesSolver,
+        is_final_step: bool | None,
+    ) -> tuple[bool, str | None, LettaStopReason | None]:
+
+        continue_stepping = request_heartbeat
+        heartbeat_reason: str | None = None
+        stop_reason: LettaStopReason | None = None
+
+        if tool_rule_violated:
+            continue_stepping = True
+            heartbeat_reason = f"{NON_USER_MSG_PREFIX}Continuing: tool rule violation."
+        else:
+            tool_rules_solver.register_tool_call(tool_call_name)
+
+            if tool_rules_solver.is_terminal_tool(tool_call_name):
+                if continue_stepping:
+                    stop_reason = LettaStopReason(stop_reason=StopReasonType.tool_rule.value)
+                continue_stepping = False
+
+            elif tool_rules_solver.has_children_tools(tool_call_name):
+                continue_stepping = True
+                heartbeat_reason = f"{NON_USER_MSG_PREFIX}Continuing: child tool rule."
+
+            elif tool_rules_solver.is_continue_tool(tool_call_name):
+                continue_stepping = True
+                heartbeat_reason = f"{NON_USER_MSG_PREFIX}Continuing: continue tool rule."
+
+        # – hard stop overrides –
+        if is_final_step:
+            continue_stepping = False
+            stop_reason = LettaStopReason(stop_reason=StopReasonType.max_steps.value)
+        else:
+            uncalled = tool_rules_solver.get_uncalled_required_tools()
+            if not continue_stepping and uncalled:
+                continue_stepping = True
+                heartbeat_reason = f"{NON_USER_MSG_PREFIX}Missing required tools: " f"{', '.join(uncalled)}"
+
+                stop_reason = None  # reset – we’re still going
+
+        return continue_stepping, heartbeat_reason, stop_reason
+
     @trace_method
     async def _execute_tool(
         self,
@@ -1102,6 +1127,7 @@ class LettaAgent(BaseAgent):
             message_manager=self.message_manager,
             agent_manager=self.agent_manager,
             block_manager=self.block_manager,
+            job_manager=self.job_manager,
             passage_manager=self.passage_manager,
             sandbox_env_vars=sandbox_env_vars,
             actor=self.actor,