Python: [Bug]: Custom Message Store State Lost During Thread Serialization

[hombin]

Description

When using custom ChatMessageStore implementations (e.g., RedisChatMessageStore from agent_framework.redis), the thread serialization/deserialization process loses custom store-specific configuration, making it impossible to restore threads from serialized state.

Environment

• Package: agent_framework
• Component: agent_framework._threads.AgentThread
• Affected Implementations: Any custom ChatMessageStoreProtocol implementation with custom fields beyond messages

Code Sample

### Steps to Reproduce

1. Create an agent with a custom message store (e.g., `RedisChatMessageStore`):


from agent_framework import ChatAgent
from agent_framework.redis import RedisChatMessageStore
from agent_framework.openai import OpenAIChatClient

client = OpenAIChatClient(model="gpt-4o", api_key="your-key")

agent = await ChatAgent(
    chat_client=client,
    name="test_agent",
    instructions="You are a helpful assistant.",
    chat_message_store_factory=lambda: RedisChatMessageStore(
        redis_url="redis://:password@127.0.0.1:6379/3",  # local redis url
        key_prefix="test_messages",
        max_messages=10,
    ),
).__aenter__()


2. Create a thread and have a conversation:


thread = agent.get_new_thread()
response = await agent.run("Hello, my name is Alice", thread=thread)


3. Serialize the thread:


serialized = await thread.serialize()
print("Serialized state:", serialized)


4. Try to deserialize and continue the conversation:


new_thread = await agent.deserialize_thread(serialized)
response = await agent.run("What's my name?", thread=new_thread)

Error Messages / Stack Traces

### Expected Behavior

The serialized state should preserve all Redis-specific configuration:


{
  'type': 'agent_thread_state',
  'service_thread_id': None,
  'chat_message_store_state': {
    'type': 'redis_store_state',
    'thread_id': 'thread_c9aee3d9-aa3a-4e08-ad14-38694fbf326e',
    'redis_url': 'redis://:password@127.0.0.1:6379/3',
    'key_prefix': 'test_messages',
    'max_messages': 10,
    'messages': []
  }
}


The deserialized thread should work correctly with the same Redis backend and be able to access the conversation history.

### Actual Behavior

The serialized state loses all Redis-specific fields:


{
  'type': 'agent_thread_state',
  'service_thread_id': None,
  'chat_message_store_state': {
    'type': 'chat_message_store_state',
    'messages': []  # ❌ Redis config lost! Only messages remain
  }
}

Package Versions

agent-framework-core[all]==1.0.0b251209

Python Version

Python 3.12.12

Additional Context

When trying to deserialize:

• The Redis configuration (redis_url, thread_id, key_prefix, max_messages) is completely lost
• Deserialization either fails or creates an in-memory store instead of reconnecting to Redis
• The thread cannot access the conversation history stored in Redis

Root Cause

In agent_framework/_threads.py, the AgentThreadState.__init__() method (around line 176) forcibly converts the chat_message_store_state dictionary to a ChatMessageStoreState object:

if isinstance(chat_message_store_state, dict):
    self.chat_message_store_state = ChatMessageStoreState.from_dict(chat_message_store_state)

The ChatMessageStoreState class only has a messages field, so all custom fields from implementations like RedisChatMessageStore (e.g., redis_url, thread_id, key_prefix, max_messages) are discarded during this conversion.

Impact

This bug affects:

• ✗ RedisChatMessageStore: Cannot persist/restore threads with Redis backend
• ✗ Any custom ChatMessageStore implementations: Cannot preserve custom configuration
• ✗ Production deployments: Thread state cannot be reliably saved/restored across service restarts
• ✗ Distributed systems: Cannot share thread state across multiple service instances

Workaround

Currently, there's no clean workaround. Users must either:

1. Use only in-memory ChatMessageStore (losing persistence benefits)
2. Manually manage serialization outside the framework
3. Re-implement thread management logic

Proposed Solution

Modify AgentThreadState.__init__() to preserve the original dictionary representation without conversion, allowing custom message store implementations to maintain their specific fields during serialization/deserialization.

The key changes needed:

1. Don't convert chat_message_store_state to ChatMessageStoreState in AgentThreadState.__init__()
2. Update AgentThread.deserialize() to handle both standard and custom message stores
3. Update AgentThread.update_from_thread_state() to handle both standard and custom message stores

This would be backward compatible since existing code using standard ChatMessageStore would continue to work.

Additional Context

This issue was discovered when implementing persistent conversation storage using Redis. The RedisChatMessageStore.serialize() method correctly returns all Redis configuration, but this information is lost when AgentThread.serialize() processes it through AgentThreadState.