feat: add last_stop_reason to AgentState [LET-5911] (#5772)

* feat: add last_stop_reason to AgentState [LET-5911]

* undo agent loop changes, use update_run_by_id_async

* add run manager test

* add integration tests

* remove comment

* remove duplicate test
This commit is contained in:
Christina Tong
2025-11-04 23:32:56 -08:00
committed by Caren Thomas
parent dbad510a6e
commit ef3df907c5
10 changed files with 248 additions and 0 deletions

View File

@@ -0,0 +1,31 @@
"""add last_stop_reason to agent state
Revision ID: 89fd4648866b
Revises: f6cd5a1e519d
Create Date: 2025-10-27 16:55:54.383688
"""
from typing import Sequence, Union
import sqlalchemy as sa
from alembic import op
# revision identifiers, used by Alembic.
revision: str = "89fd4648866b"
down_revision: Union[str, None] = "f6cd5a1e519d"
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.add_column("agents", sa.Column("last_stop_reason", sa.String(), nullable=True))
# ### end Alembic commands ###
def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ###
op.drop_column("agents", "last_stop_reason")
# ### end Alembic commands ###

View File

@@ -18981,6 +18981,17 @@
"title": "Last Run Duration Ms",
"description": "The duration in milliseconds of the agent's last run."
},
"last_stop_reason": {
"anyOf": [
{
"$ref": "#/components/schemas/StopReasonType"
},
{
"type": "null"
}
],
"description": "The stop reason from the agent's last run."
},
"timezone": {
"anyOf": [
{
@@ -34653,6 +34664,17 @@
"title": "Last Run Duration Ms",
"description": "The duration in milliseconds of the agent's last run."
},
"last_stop_reason": {
"anyOf": [
{
"$ref": "#/components/schemas/StopReasonType"
},
{
"type": "null"
}
],
"description": "The stop reason from the agent's last run."
},
"timezone": {
"anyOf": [
{

View File

@@ -16,6 +16,7 @@ from letta.orm.sqlalchemy_base import SqlalchemyBase
from letta.schemas.agent import AgentState as PydanticAgentState
from letta.schemas.embedding_config import EmbeddingConfig
from letta.schemas.enums import AgentType
from letta.schemas.letta_stop_reason import StopReasonType
from letta.schemas.llm_config import LLMConfig
from letta.schemas.memory import Memory
from letta.schemas.response_format import ResponseFormatUnion
@@ -93,6 +94,9 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
last_run_duration_ms: Mapped[Optional[int]] = mapped_column(
Integer, nullable=True, doc="The duration in milliseconds of the agent's last run."
)
last_stop_reason: Mapped[Optional[StopReasonType]] = mapped_column(
String, nullable=True, doc="The stop reason from the agent's last run."
)
# timezone
timezone: Mapped[Optional[str]] = mapped_column(String, nullable=True, doc="The timezone of the agent (for the context window).")
@@ -232,6 +236,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
"response_format": self.response_format,
"last_run_completion": self.last_run_completion,
"last_run_duration_ms": self.last_run_duration_ms,
"last_stop_reason": self.last_stop_reason,
"timezone": self.timezone,
"max_files_open": self.max_files_open,
"per_file_view_window_char_limit": self.per_file_view_window_char_limit,
@@ -334,6 +339,7 @@ class Agent(SqlalchemyBase, OrganizationMixin, ProjectMixin, TemplateEntityMixin
"response_format": self.response_format,
"last_run_completion": self.last_run_completion,
"last_run_duration_ms": self.last_run_duration_ms,
"last_stop_reason": self.last_stop_reason,
"max_files_open": self.max_files_open,
"per_file_view_window_char_limit": self.per_file_view_window_char_limit,
"hidden": self.hidden,

View File

@@ -14,6 +14,7 @@ from letta.schemas.file import FileStatus
from letta.schemas.group import Group
from letta.schemas.identity import Identity
from letta.schemas.letta_base import OrmMetadataBase
from letta.schemas.letta_stop_reason import StopReasonType
from letta.schemas.llm_config import LLMConfig
from letta.schemas.memory import Memory
from letta.schemas.message import Message, MessageCreate
@@ -135,6 +136,7 @@ class AgentState(OrmMetadataBase, validate_assignment=True):
# Run metrics
last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.")
last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.")
last_stop_reason: Optional[StopReasonType] = Field(None, description="The stop reason from the agent's last run.")
# timezone
timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).")
@@ -388,6 +390,7 @@ class UpdateAgent(BaseModel):
response_format: Optional[ResponseFormatUnion] = Field(None, description="The response format for the agent.")
last_run_completion: Optional[datetime] = Field(None, description="The timestamp when the agent last completed a run.")
last_run_duration_ms: Optional[int] = Field(None, description="The duration in milliseconds of the agent's last run.")
last_stop_reason: Optional[StopReasonType] = Field(None, description="The stop reason from the agent's last run.")
timezone: Optional[str] = Field(None, description="The timezone of the agent (IANA format).")
max_files_open: Optional[int] = Field(
None,

View File

@@ -712,6 +712,7 @@ class AgentManager:
"response_format": agent_update.response_format,
"last_run_completion": agent_update.last_run_completion,
"last_run_duration_ms": agent_update.last_run_duration_ms,
"last_stop_reason": agent_update.last_stop_reason,
"timezone": agent_update.timezone,
"max_files_open": agent_update.max_files_open,
"per_file_view_window_char_limit": agent_update.per_file_view_window_char_limit,

View File

@@ -341,6 +341,20 @@ class RunManager:
await session.commit()
# Update agent's last_stop_reason when run completes
# Do this after run update is committed to database
if is_terminal_update and update.stop_reason:
try:
from letta.schemas.agent import UpdateAgent
await self.agent_manager.update_agent_async(
agent_id=pydantic_run.agent_id,
agent_update=UpdateAgent(last_stop_reason=update.stop_reason),
actor=actor,
)
except Exception as e:
logger.error(f"Failed to update agent's last_stop_reason for run {run_id}: {e}")
# update run metrics table
num_steps = len(await self.step_manager.list_steps_async(run_id=run_id, actor=actor))

View File

@@ -1277,3 +1277,51 @@ def test_parallel_tool_calling(
assert messages[1].message_type == "assistant_message"
assert messages[2].message_type == "stop_reason"
assert messages[3].message_type == "usage_statistics"
def test_agent_records_last_stop_reason_after_approval_flow(
client: Letta,
agent: AgentState,
) -> None:
"""
Test that the agent's last_stop_reason is properly updated after a human-in-the-loop flow.
This verifies the integration between run completion and agent state updates.
"""
# Get initial agent state
initial_agent = client.agents.retrieve(agent_id=agent.id)
initial_stop_reason = initial_agent.last_stop_reason
# Trigger approval request
response = client.agents.messages.create(
agent_id=agent.id,
messages=USER_MESSAGE_TEST_APPROVAL,
)
# Verify we got an approval request
messages = response.messages
assert messages is not None
assert len(messages) == 3
assert messages[2].message_type == "approval_request_message"
# Check agent after approval request (run should be paused with requires_approval)
agent_after_request = client.agents.retrieve(agent_id=agent.id)
assert agent_after_request.last_stop_reason == "requires_approval"
# Approve the tool call
approve_tool_call(client, agent.id, response.messages[2].tool_call.tool_call_id)
# Check agent after approval (run should complete with end_turn or similar)
agent_after_approval = client.agents.retrieve(agent_id=agent.id)
# After approval and run completion, stop reason should be updated (could be end_turn or other terminal reason)
assert agent_after_approval.last_stop_reason is not None
assert agent_after_approval.last_stop_reason != initial_stop_reason # Should be different from initial
# Send follow-up message to complete the flow
response2 = client.agents.messages.create(
agent_id=agent.id,
messages=USER_MESSAGE_FOLLOW_UP,
)
# Verify final agent state has the most recent stop reason
final_agent = client.agents.retrieve(agent_id=agent.id)
assert final_agent.last_stop_reason is not None

View File

@@ -625,6 +625,46 @@ async def test_update_agent_file_fields(server: SyncServer, comprehensive_test_a
assert updated_agent.per_file_view_window_char_limit == 150_000
@pytest.mark.asyncio
async def test_update_agent_last_stop_reason(server: SyncServer, comprehensive_test_agent_fixture, default_user):
"""Test updating last_stop_reason field on an existing agent"""
agent, _ = comprehensive_test_agent_fixture
assert agent.last_stop_reason is None
# Update with end_turn stop reason
update_request = UpdateAgent(
last_stop_reason=StopReasonType.end_turn,
last_run_completion=datetime.now(timezone.utc),
last_run_duration_ms=1500,
)
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
assert updated_agent.last_stop_reason == StopReasonType.end_turn
assert updated_agent.last_run_completion is not None
assert updated_agent.last_run_duration_ms == 1500
# Update with error stop reason
update_request = UpdateAgent(
last_stop_reason=StopReasonType.error,
last_run_completion=datetime.now(timezone.utc),
last_run_duration_ms=2500,
)
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
assert updated_agent.last_stop_reason == StopReasonType.error
assert updated_agent.last_run_duration_ms == 2500
# Update with requires_approval stop reason
update_request = UpdateAgent(
last_stop_reason=StopReasonType.requires_approval,
)
updated_agent = await server.agent_manager.update_agent_async(agent.id, update_request, actor=default_user)
assert updated_agent.last_stop_reason == StopReasonType.requires_approval
# ======================================================================================================================
# AgentManager Tests - Listing
# ======================================================================================================================
@@ -1086,6 +1126,7 @@ async def test_agent_state_schema_unchanged(server: SyncServer):
# Run metrics
"last_run_completion": (datetime, type(None)),
"last_run_duration_ms": (int, type(None)),
"last_stop_reason": (StopReasonType, type(None)),
# Timezone
"timezone": (str, type(None)),
# File controls

View File

@@ -232,6 +232,41 @@ async def test_update_run_metadata_persistence(server: SyncServer, sarah_agent,
assert fetched_run.metadata["error"]["type"] == "llm_timeout"
@pytest.mark.asyncio
async def test_update_run_updates_agent_last_stop_reason(server: SyncServer, sarah_agent, default_user):
"""Test that completing a run updates the agent's last_stop_reason."""
# Verify agent starts with no last_stop_reason
agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
initial_stop_reason = agent.last_stop_reason
# Create a run
run_data = PydanticRun(agent_id=sarah_agent.id)
created_run = await server.run_manager.create_run(pydantic_run=run_data, actor=default_user)
# Complete the run with end_turn stop reason
await server.run_manager.update_run_by_id_async(
created_run.id, RunUpdate(status=RunStatus.completed, stop_reason=StopReasonType.end_turn), actor=default_user
)
# Verify agent's last_stop_reason was updated to end_turn
updated_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
assert updated_agent.last_stop_reason == StopReasonType.end_turn
# Create another run and complete with different stop reason
run_data2 = PydanticRun(agent_id=sarah_agent.id)
created_run2 = await server.run_manager.create_run(pydantic_run=run_data2, actor=default_user)
# Complete with error stop reason
await server.run_manager.update_run_by_id_async(
created_run2.id, RunUpdate(status=RunStatus.failed, stop_reason=StopReasonType.error), actor=default_user
)
# Verify agent's last_stop_reason was updated to error
final_agent = await server.agent_manager.get_agent_by_id_async(agent_id=sarah_agent.id, actor=default_user)
assert final_agent.last_stop_reason == StopReasonType.error
@pytest.mark.asyncio
async def test_delete_run_by_id(server: SyncServer, sarah_agent, default_user):
"""Test deleting a run by its ID."""

View File

@@ -1179,3 +1179,50 @@ def test_parallel_tool_calling(
assert messages[1].message_type == "assistant_message"
assert messages[2].message_type == "stop_reason"
assert messages[3].message_type == "usage_statistics"
def test_agent_records_last_stop_reason_after_approval_flow(
client: Letta,
agent: AgentState,
) -> None:
"""
Test that the agent's last_stop_reason is properly updated after a human-in-the-loop flow.
This verifies the integration between run completion and agent state updates.
"""
# Get initial agent state
initial_agent = client.agents.retrieve(agent_id=agent.id)
initial_stop_reason = initial_agent.last_stop_reason
# Trigger approval request
response = client.agents.messages.send(
agent_id=agent.id,
messages=USER_MESSAGE_TEST_APPROVAL,
)
# Verify we got an approval request
messages = response.messages
assert messages is not None
assert len(messages) == 3
assert messages[2].message_type == "approval_request_message"
# Check agent after approval request (run should be paused with requires_approval)
agent_after_request = client.agents.retrieve(agent_id=agent.id)
assert agent_after_request.last_stop_reason == "requires_approval"
# Approve the tool call
approve_tool_call(client, agent.id, response.messages[2].tool_call.tool_call_id)
# Check agent after approval (run should complete with end_turn or similar)
agent_after_approval = client.agents.retrieve(agent_id=agent.id)
assert agent_after_approval.last_stop_reason is not None
assert agent_after_approval.last_stop_reason != initial_stop_reason
# Send follow-up message to complete the flow
response2 = client.agents.messages.send(
agent_id=agent.id,
messages=USER_MESSAGE_FOLLOW_UP,
)
# Verify final agent state has the most recent stop reason
final_agent = client.agents.retrieve(agent_id=agent.id)
assert final_agent.last_stop_reason is not None