From 66047d54f779880eabcc375c25ba1e802b23b33c Mon Sep 17 00:00:00 2001 From: Matthew Zhou Date: Thu, 24 Jul 2025 15:14:00 -0700 Subject: [PATCH] fix: Fix test managers state transition tests (#3551) --- letta/services/file_manager.py | 42 +++++++++++++++++++--------------- tests/test_managers.py | 1 + 2 files changed, 24 insertions(+), 19 deletions(-) diff --git a/letta/services/file_manager.py b/letta/services/file_manager.py index 5dbbe5d7..6dc06778 100644 --- a/letta/services/file_manager.py +++ b/letta/services/file_manager.py @@ -151,7 +151,8 @@ class FileManager: Enforces state transition rules (when enforce_state_transitions=True): - PENDING -> PARSING -> EMBEDDING -> COMPLETED (normal flow) - Any non-terminal state -> ERROR - - ERROR and COMPLETED are terminal (no transitions allowed) + - Same-state transitions are allowed (e.g., EMBEDDING -> EMBEDDING) + - ERROR and COMPLETED are terminal (no status transitions allowed, metadata updates blocked) Args: file_id: ID of the file to update @@ -196,28 +197,31 @@ class FileManager: ] # only add state transition validation if enforce_state_transitions is True - if enforce_state_transitions: - # prevent updates to terminal states (ERROR, COMPLETED) + if enforce_state_transitions and processing_status is not None: + # enforce specific transitions based on target status + if processing_status == FileProcessingStatus.PARSING: + where_conditions.append( + FileMetadataModel.processing_status.in_([FileProcessingStatus.PENDING, FileProcessingStatus.PARSING]) + ) + elif processing_status == FileProcessingStatus.EMBEDDING: + where_conditions.append( + FileMetadataModel.processing_status.in_([FileProcessingStatus.PARSING, FileProcessingStatus.EMBEDDING]) + ) + elif processing_status == FileProcessingStatus.COMPLETED: + where_conditions.append( + FileMetadataModel.processing_status.in_([FileProcessingStatus.EMBEDDING, FileProcessingStatus.COMPLETED]) + ) + elif processing_status == FileProcessingStatus.ERROR: + # ERROR can be set from any non-terminal state + where_conditions.append( + FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED]) + ) + elif enforce_state_transitions and processing_status is None: + # If only updating metadata fields (not status), prevent updates to terminal states where_conditions.append( FileMetadataModel.processing_status.notin_([FileProcessingStatus.ERROR, FileProcessingStatus.COMPLETED]) ) - if processing_status is not None: - # enforce specific transitions based on target status - if processing_status == FileProcessingStatus.PARSING: - where_conditions.append( - FileMetadataModel.processing_status.in_([FileProcessingStatus.PENDING, FileProcessingStatus.PARSING]) - ) - elif processing_status == FileProcessingStatus.EMBEDDING: - where_conditions.append( - FileMetadataModel.processing_status.in_([FileProcessingStatus.PARSING, FileProcessingStatus.EMBEDDING]) - ) - elif processing_status == FileProcessingStatus.COMPLETED: - where_conditions.append( - FileMetadataModel.processing_status.in_([FileProcessingStatus.EMBEDDING, FileProcessingStatus.COMPLETED]) - ) - # ERROR can be set from any non-terminal state (already handled by terminal check above) - # fast in-place update with state validation stmt = ( update(FileMetadataModel) diff --git a/tests/test_managers.py b/tests/test_managers.py index 2556b9d4..c75b6167 100644 --- a/tests/test_managers.py +++ b/tests/test_managers.py @@ -55,6 +55,7 @@ from letta.schemas.block import BlockUpdate, CreateBlock from letta.schemas.embedding_config import EmbeddingConfig from letta.schemas.enums import ActorType, AgentStepStatus, FileProcessingStatus, JobStatus, JobType, MessageRole, ProviderType from letta.schemas.environment_variables import SandboxEnvironmentVariableCreate, SandboxEnvironmentVariableUpdate +from letta.schemas.file import FileMetadata from letta.schemas.file import FileMetadata as PydanticFileMetadata from letta.schemas.identity import IdentityCreate, IdentityProperty, IdentityPropertyType, IdentityType, IdentityUpdate, IdentityUpsert from letta.schemas.job import BatchJob