From 1d58bbe7519d23afb8c294f566d8f8f675510fcc Mon Sep 17 00:00:00 2001 From: nanxingw Date: Mon, 27 Apr 2026 16:35:40 +0800 Subject: [PATCH] [VEPEAGE-000] Add top-level session_id field to messages Introduce a session_id column on messages so callers can group messages into logical conversation sessions independent of agent_id. The column is optional (nullable), validated at the application layer (regex [A-Za-z0-9_-]+, max length 64) and backed by a Postgres CHECK constraint plus a covering index (agent_id, session_id, created_at). Touch points: - ORM/schema: mirix/orm/message.py adds the column, ix_messages_agent_session_created_at index, and ck_messages_session_id_format check (Postgres-only); shared pattern/length constants live in mirix/schemas/message.py. - Validation: MessageCreate / MessageUpdate / Message all enforce the same rule via a shared _validate_session_id helper. - Queue: message.proto gets field 8 (session_id, optional string); message_pb2.py / message_pb2.pyi regenerated with grpcio-tools (protobuf floor bumped to 5.27.2 in pyproject.toml + requirements.txt to match the gencode header). New `make proto` target keeps regen reproducible. - Plumbing: queue_util / worker propagate session_id end-to-end; message_helpers.prepare_input_message_create forwards it; agent.py inherits the triggering input's session_id onto every synthesized assistant/tool/heartbeat/summary message in the same step (via a step-level _current_step_session_id stash that step_user_message save/restores so summarize_messages_inplace gets the right value). - API: SendMessageRequest and AddMemoryRequest gain session_id with the same validator; AddMemoryRequest enforces agreement between top-level session_id and filter_tags["session_id"] when both are set. - Filtering: MessageManager.list_messages_for_{agent,user} accept session_id and translate to a column-equality filter. - Migrations: scripts/migrate_add_message_session_id.sql adds the column + index + check; _phase2.sql is the follow-up backfill/tighten step. - Tests: tests/test_session_id.py (DB-free unit, ~470 lines covering schema, validator, helper, queue serialization, ORM column shape) and tests/test_session_id_integration.py (Postgres round-trip via MessageManager confirming session-scoped filtering). Co-Authored-By: Claude Opus 4.7 (1M context) --- Makefile | 13 +- mirix/agent/agent.py | 388 ++++++++------ mirix/helpers/message_helpers.py | 1 + mirix/orm/message.py | 32 +- mirix/queue/message.proto | 3 + mirix/queue/message_pb2.py | 32 +- mirix/queue/message_pb2.pyi | 6 +- mirix/queue/queue_util.py | 2 + mirix/queue/worker.py | 1 + mirix/schemas/message.py | 79 +++ mirix/server/rest_api.py | 53 +- mirix/services/message_manager.py | 6 + pyproject.toml | 2 +- requirements.txt | 2 +- scripts/migrate_add_message_session_id.sql | 54 ++ .../migrate_add_message_session_id_phase2.sql | 11 + tests/test_session_id.py | 497 ++++++++++++++++++ tests/test_session_id_integration.py | 213 ++++++++ 18 files changed, 1205 insertions(+), 190 deletions(-) create mode 100644 scripts/migrate_add_message_session_id.sql create mode 100644 scripts/migrate_add_message_session_id_phase2.sql create mode 100644 tests/test_session_id.py create mode 100644 tests/test_session_id_integration.py diff --git a/Makefile b/Makefile index 942f69301..d348102cf 100644 --- a/Makefile +++ b/Makefile @@ -1,4 +1,4 @@ -.PHONY: install format lint test all check +.PHONY: install format lint test all check proto # Define variables PYTHON = python3 @@ -30,3 +30,14 @@ test: # Run format, lint, and test check: format lint test + +# Regenerate protobuf gencode for mirix/queue/*.proto. +# Uses grpcio-tools pinned in pyproject.toml (>=1.66.0,<1.67.0) so the +# checked-in *_pb2.py / *_pb2.pyi / *_pb2_grpc.py files are reproducible. +proto: + $(POETRY) run python -m grpc_tools.protoc \ + -I. \ + --python_out=. \ + --pyi_out=. \ + --grpc_python_out=. \ + mirix/queue/message.proto diff --git a/mirix/agent/agent.py b/mirix/agent/agent.py index 953dbeaf8..ab65c45e5 100644 --- a/mirix/agent/agent.py +++ b/mirix/agent/agent.py @@ -946,6 +946,10 @@ async def _handle_ai_response( if response_message_id is not None: assert response_message_id.startswith("message-"), response_message_id + # Inherit the triggering input's session_id so synthesized assistant/tool + # messages from this step stay in the same session (Codex review C1). + input_session_id = getattr(input_message, "session_id", None) + messages = [] # append these to the history when done function_name = None @@ -988,6 +992,7 @@ async def _handle_ai_response( agent_id=self.agent_state.id, model=self.model, openai_message_dict=response_message.model_dump(), + session_id=input_session_id, ) ) # extend conversation with assistant's reply @@ -1037,6 +1042,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Error: {error_msg}", msg_obj=messages[-1]) @@ -1062,6 +1068,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Error: {error_msg}", msg_obj=messages[-1]) @@ -1111,6 +1118,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) self.interface.function_message(f"Validation Error: {validation_error}", msg_obj=messages[-1]) @@ -1197,6 +1205,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1218,6 +1227,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1237,6 +1247,7 @@ async def _handle_ai_response( "content": function_response, "tool_call_id": tool_call_id, }, + session_id=input_session_id, ) ) # extend conversation with function response self.interface.function_message(f"Ran {function_name}()", msg_obj=messages[-1]) @@ -1423,6 +1434,7 @@ async def _handle_ai_response( "role": "user", "content": message_content, }, + session_id=input_session_id, ) # persist the message to the database @@ -1465,6 +1477,7 @@ async def _handle_ai_response( agent_id=self.agent_state.id, model=self.model, openai_message_dict=response_message.model_dump(), + session_id=input_session_id, ) ) # extend conversation with assistant's reply self.interface.internal_monologue(response_message.content, msg_obj=messages[-1]) @@ -1541,187 +1554,205 @@ async def step( first_input_message = input_messages[0] if isinstance(input_messages, list) else input_messages - # Convert MessageCreate objects to Message objects - if not isinstance(input_messages, list): - input_messages = [input_messages] - message_objects = [ - ( - m - if isinstance(m, Message) - else prepare_input_message_create( - m, - self.agent_state.id, - wrap_user_message=False, - wrap_system_message=True, - ) - ) - for m in input_messages - ] + # Derive the session_id for this step once so all synthesized messages + # (heartbeats, warnings, meta-memory bootstrap, summaries) inherit it. + step_session_id = getattr(first_input_message, "session_id", None) + # Expose on self so helpers called without an explicit session_id + # (e.g. summarize_messages_inplace) can pick it up. Save the prior + # value and restore in `finally` so a step() that raises or completes + # cannot leave a stale session_id stashed for the next caller (Codex review v4). + prev_session_id = getattr(self, "_current_step_session_id", None) + self._current_step_session_id = step_session_id + try: - extra_message_objects = ( - [ - prepare_input_message_create( - m, - self.agent_state.id, - wrap_user_message=False, - wrap_system_message=True, + # Convert MessageCreate objects to Message objects + if not isinstance(input_messages, list): + input_messages = [input_messages] + message_objects = [ + ( + m + if isinstance(m, Message) + else prepare_input_message_create( + m, + self.agent_state.id, + wrap_user_message=False, + wrap_system_message=True, + ) ) - for m in extra_messages + for m in input_messages ] - if extra_messages is not None - else None - ) - next_input_message = message_objects - counter = 0 - total_usage = UsageStatistics() - step_count = 0 - initial_message_count = len( - await self.agent_manager.get_in_context_messages( - agent_state=self.agent_state, actor=self.actor, user=self.user + extra_message_objects = ( + [ + prepare_input_message_create( + m, + self.agent_state.id, + wrap_user_message=False, + wrap_system_message=True, + ) + for m in extra_messages + ] + if extra_messages is not None + else None ) - ) + next_input_message = message_objects + counter = 0 + total_usage = UsageStatistics() + step_count = 0 - if self.agent_state.is_type(AgentType.reflexion_agent): - # clear previous messages - in_context_messages = await self.agent_manager.get_in_context_messages( - agent_state=self.agent_state, actor=self.actor, user=self.user - ) - in_context_messages = in_context_messages[:1] - await self.agent_manager.set_in_context_messages( - agent_id=self.agent_state.id, - message_ids=[message.id for message in in_context_messages], - actor=self.actor, + initial_message_count = len( + await self.agent_manager.get_in_context_messages( + agent_state=self.agent_state, actor=self.actor, user=self.user + ) ) - # Initialize the LLM client once per step to reuse across retries. - llm_client = LLMClient.create( - llm_config=self.agent_state.llm_config, - ) + if self.agent_state.is_type(AgentType.reflexion_agent): + # clear previous messages + in_context_messages = await self.agent_manager.get_in_context_messages( + agent_state=self.agent_state, actor=self.actor, user=self.user + ) + in_context_messages = in_context_messages[:1] + await self.agent_manager.set_in_context_messages( + agent_id=self.agent_state.id, + message_ids=[message.id for message in in_context_messages], + actor=self.actor, + ) - while True: - kwargs["first_message"] = False - kwargs["step_count"] = step_count + # Initialize the LLM client once per step to reuse across retries. + llm_client = LLMClient.create( + llm_config=self.agent_state.llm_config, + ) - if self.agent_state.is_type(AgentType.meta_memory_agent, AgentType.chat_agent) and step_count == 0: - # When the agent first gets the screenshots, we need to extract the topic to search the query. - try: - topics = await self._extract_topics_from_messages(next_input_message) + while True: + kwargs["first_message"] = False + kwargs["step_count"] = step_count - if topics is not None: - kwargs["topics"] = topics - else: - printv(f"[Mirix.Agent.{self.agent_state.name}] WARNING: No topics extracted from screenshots") + if self.agent_state.is_type(AgentType.meta_memory_agent, AgentType.chat_agent) and step_count == 0: + # When the agent first gets the screenshots, we need to extract the topic to search the query. + try: + topics = await self._extract_topics_from_messages(next_input_message) - except Exception as e: - printv( - f"[Mirix.Agent.{self.agent_state.name}] INFO: Error in extracting the topic from the screenshots: {e}" + if topics is not None: + kwargs["topics"] = topics + else: + printv(f"[Mirix.Agent.{self.agent_state.name}] WARNING: No topics extracted from screenshots") + + except Exception as e: + printv( + f"[Mirix.Agent.{self.agent_state.name}] INFO: Error in extracting the topic from the screenshots: {e}" + ) + pass + + if self.agent_state.is_type(AgentType.meta_memory_agent) and step_count == 0: + meta_message = prepare_input_message_create( + MessageCreate( + role="user", + content="[System Message] As the meta memory manager, analyze the provided content and perform your function.", + filter_tags=self.filter_tags, + session_id=step_session_id, + ), + self.agent_state.id, + wrap_user_message=False, + wrap_system_message=True, ) - pass + next_input_message.append(meta_message) - if self.agent_state.is_type(AgentType.meta_memory_agent) and step_count == 0: - meta_message = prepare_input_message_create( - MessageCreate( - role="user", - content="[System Message] As the meta memory manager, analyze the provided content and perform your function.", - filter_tags=self.filter_tags, - ), - self.agent_state.id, - wrap_user_message=False, - wrap_system_message=True, + step_response = await self.inner_step( + first_input_messge=first_input_message, + messages=next_input_message, + extra_messages=extra_message_objects, + initial_message_count=initial_message_count, + chaining=chaining, + llm_client=llm_client, + **kwargs, ) - next_input_message.append(meta_message) - - step_response = await self.inner_step( - first_input_messge=first_input_message, - messages=next_input_message, - extra_messages=extra_message_objects, - initial_message_count=initial_message_count, - chaining=chaining, - llm_client=llm_client, - **kwargs, - ) - continue_chaining = step_response.continue_chaining - function_failed = step_response.function_failed - token_warning = step_response.in_context_memory_warning - usage = step_response.usage + continue_chaining = step_response.continue_chaining + function_failed = step_response.function_failed + token_warning = step_response.in_context_memory_warning + usage = step_response.usage - step_count += 1 - total_usage += usage - counter += 1 - self.interface.step_complete() + step_count += 1 + total_usage += usage + counter += 1 + self.interface.step_complete() - # logger.debug("Saving agent state") - # save updated state - await save_agent(self) + # logger.debug("Saving agent state") + # save updated state + await save_agent(self) - # Chain stops - if not chaining and (not function_failed): - printv(f"[Mirix.Agent.{self.agent_state.name}] INFO: No chaining, stopping after one step") - break - elif max_chaining_steps is not None and counter == max_chaining_steps: - # Add warning message based on agent type - if self.agent_state.is_type(AgentType.chat_agent): - warning_content = "[System Message] You have reached the maximum chaining steps. Please call 'send_message' to send your response to the user." + # Chain stops + if not chaining and (not function_failed): + printv(f"[Mirix.Agent.{self.agent_state.name}] INFO: No chaining, stopping after one step") + break + elif max_chaining_steps is not None and counter == max_chaining_steps: + # Add warning message based on agent type + if self.agent_state.is_type(AgentType.chat_agent): + warning_content = "[System Message] You have reached the maximum chaining steps. Please call 'send_message' to send your response to the user." + else: + warning_content = "[System Message] You have reached the maximum chaining steps. Please call 'finish_memory_update' to end the chaining." + next_input_message = Message.dict_to_message( + agent_id=self.agent_state.id, + model=self.model, + openai_message_dict={ + "role": "user", + "content": warning_content, + }, + session_id=step_session_id, + ) + continue # give agent one more chance to respond + elif max_chaining_steps is not None and counter > max_chaining_steps: + printv( + f"[Mirix.Agent.{self.agent_state.name}] INFO: Hit max chaining steps, stopping after {counter} steps" + ) + break + # Chain handlers + elif token_warning and summarizer_settings.send_memory_warning_message: + assert self.agent_state.created_by_id is not None + next_input_message = Message.dict_to_message( + agent_id=self.agent_state.id, + model=self.model, + openai_message_dict={ + "role": "user", # TODO: change to system? + "content": get_token_limit_warning(), + }, + session_id=step_session_id, + ) + continue # always chain + elif function_failed: + assert self.agent_state.created_by_id is not None + next_input_message = Message.dict_to_message( + agent_id=self.agent_state.id, + model=self.model, + openai_message_dict={ + "role": "user", # TODO: change to system? + "content": get_contine_chaining(FUNC_FAILED_HEARTBEAT_MESSAGE), + }, + session_id=step_session_id, + ) + continue # always chain + elif continue_chaining: + assert self.agent_state.created_by_id is not None + next_input_message = Message.dict_to_message( + agent_id=self.agent_state.id, + model=self.model, + openai_message_dict={ + "role": "user", # TODO: change to system? + "content": get_contine_chaining(REQ_HEARTBEAT_MESSAGE), + }, + session_id=step_session_id, + ) + continue # always chain + # Mirix no-op / yield else: - warning_content = "[System Message] You have reached the maximum chaining steps. Please call 'finish_memory_update' to end the chaining." - next_input_message = Message.dict_to_message( - agent_id=self.agent_state.id, - model=self.model, - openai_message_dict={ - "role": "user", - "content": warning_content, - }, - ) - continue # give agent one more chance to respond - elif max_chaining_steps is not None and counter > max_chaining_steps: - printv( - f"[Mirix.Agent.{self.agent_state.name}] INFO: Hit max chaining steps, stopping after {counter} steps" - ) - break - # Chain handlers - elif token_warning and summarizer_settings.send_memory_warning_message: - assert self.agent_state.created_by_id is not None - next_input_message = Message.dict_to_message( - agent_id=self.agent_state.id, - model=self.model, - openai_message_dict={ - "role": "user", # TODO: change to system? - "content": get_token_limit_warning(), - }, - ) - continue # always chain - elif function_failed: - assert self.agent_state.created_by_id is not None - next_input_message = Message.dict_to_message( - agent_id=self.agent_state.id, - model=self.model, - openai_message_dict={ - "role": "user", # TODO: change to system? - "content": get_contine_chaining(FUNC_FAILED_HEARTBEAT_MESSAGE), - }, - ) - continue # always chain - elif continue_chaining: - assert self.agent_state.created_by_id is not None - next_input_message = Message.dict_to_message( - agent_id=self.agent_state.id, - model=self.model, - openai_message_dict={ - "role": "user", # TODO: change to system? - "content": get_contine_chaining(REQ_HEARTBEAT_MESSAGE), - }, - ) - continue # always chain - # Mirix no-op / yield - else: - break + break - # Save the message_ids - await save_agent(self) + # Save the message_ids + await save_agent(self) - return MirixUsageStatistics(**total_usage.model_dump(), step_count=step_count) + return MirixUsageStatistics(**total_usage.model_dump(), step_count=step_count) + finally: + self._current_step_session_id = prev_session_id async def build_system_prompt_with_memories( self, @@ -2561,7 +2592,12 @@ async def inner_step( ) raise e - async def step_user_message(self, user_message_str: str, **kwargs) -> AgentStepResponse: + async def step_user_message( + self, + user_message_str: str, + session_id: Optional[str] = None, + **kwargs, + ) -> AgentStepResponse: """Takes a basic user message string, turns it into a stringified JSON with extra metadata, then sends it to the agent Example: @@ -2593,12 +2629,25 @@ async def step_user_message(self, user_message_str: str, **kwargs) -> AgentStepR agent_id=self.agent_state.id, model=self.model, openai_message_dict=openai_message_dict, + session_id=session_id, # created_at=timestamp, ) - return await self.inner_step(messages=[user_message], **kwargs) + # Seed the step-level session context so any pre-persist summarization + # triggered inside inner_step() / _get_ai_reply() stamps the correct + # session_id on the summary message (Codex review v3, Important). + prev_session_id = getattr(self, "_current_step_session_id", None) + self._current_step_session_id = session_id + try: + return await self.inner_step(messages=[user_message], **kwargs) + finally: + self._current_step_session_id = prev_session_id - async def summarize_messages_inplace(self, existing_file_uris: Optional[List[str]] = None): + async def summarize_messages_inplace( + self, + existing_file_uris: Optional[List[str]] = None, + session_id: Optional[str] = None, + ): in_context_messages = await self.agent_manager.get_in_context_messages( agent_state=self.agent_state, actor=self.actor, user=self.user ) @@ -2674,13 +2723,26 @@ async def summarize_messages_inplace(self, existing_file_uris: Optional[List[str ) packed_summary_message = {"role": "user", "content": summary_message} - # Prepend the summary + # Prepend the summary. Preference order for the summary's session_id: + # 1. explicit `session_id` argument (caller-provided, e.g. step_user_message) + # 2. _current_step_session_id stashed by Agent.step() for the current step + # 3. the latest in-context message's session_id (inherits whatever session + # the conversation being summarized is already in) + summary_session_id = session_id + if summary_session_id is None: + summary_session_id = getattr(self, "_current_step_session_id", None) + if summary_session_id is None: + for m in reversed(in_context_messages): + if getattr(m, "session_id", None): + summary_session_id = m.session_id + break self.agent_state = await self.agent_manager.prepend_to_in_context_messages( messages=[ Message.dict_to_message( agent_id=self.agent_state.id, model=self.model, openai_message_dict=packed_summary_message, + session_id=summary_session_id, ) ], agent_id=self.agent_state.id, diff --git a/mirix/helpers/message_helpers.py b/mirix/helpers/message_helpers.py index 631c81622..46e03be94 100644 --- a/mirix/helpers/message_helpers.py +++ b/mirix/helpers/message_helpers.py @@ -62,4 +62,5 @@ def prepare_input_message_create(message: MessageCreate, agent_id: str, **kwargs otid=message.otid, sender_id=message.sender_id, group_id=message.group_id, + session_id=message.session_id, ) diff --git a/mirix/orm/message.py b/mirix/orm/message.py index bab3fa91b..ef9bd944b 100755 --- a/mirix/orm/message.py +++ b/mirix/orm/message.py @@ -1,6 +1,6 @@ from typing import TYPE_CHECKING, List, Optional -from sqlalchemy import JSON, ForeignKey, Index +from sqlalchemy import JSON, CheckConstraint, ForeignKey, Index, String from sqlalchemy.orm import Mapped, declared_attr, mapped_column, relationship from mirix.orm.custom_columns import ( @@ -10,8 +10,12 @@ ) from mirix.orm.mixins import AgentMixin, OrganizationMixin, UserMixin from mirix.orm.sqlalchemy_base import SqlalchemyBase -from mirix.schemas.message import Message as PydanticMessage -from mirix.schemas.message import ToolReturn +from mirix.schemas.message import ( + SESSION_ID_MAX_LEN, + SESSION_ID_SQL_PATTERN, + Message as PydanticMessage, + ToolReturn, +) from mirix.schemas.mirix_message_content import MessageContent from mirix.schemas.mirix_message_content import TextContent as PydanticTextContent from mirix.schemas.openai.openai import ToolCall as OpenAIToolCall @@ -32,6 +36,21 @@ class Message(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): Index("ix_messages_created_at", "created_at", "id"), Index("ix_messages_client_user", "client_id", "user_id"), Index("ix_messages_agent_client_user", "agent_id", "client_id", "user_id"), + # Accelerates "list messages of an agent in a session, newest first". + Index( + "ix_messages_agent_session_created_at", + "agent_id", + "session_id", + "created_at", + ), + # Backstop the app-level validator: DB must never store an invalid session_id. + # Uses the Postgres `~` operator, so emit the constraint only on Postgres + # (SQLite, used for some local/test setups, has no POSIX regex operator). + # Pattern derived from mirix.schemas.message so there's one source of truth. + CheckConstraint( + f"session_id IS NULL OR session_id ~ '{SESSION_ID_SQL_PATTERN}'", + name="ck_messages_session_id_format", + ).ddl_if(dialect="postgresql"), ) __pydantic_model__ = PydanticMessage @@ -78,6 +97,13 @@ class Message(SqlalchemyBase, OrganizationMixin, UserMixin, AgentMixin): nullable=True, doc="The id of the sender of the message, can be an identity id or agent id", ) + session_id: Mapped[Optional[str]] = mapped_column( + String(SESSION_ID_MAX_LEN), + nullable=True, + doc="Top-level conversation/session identifier for grouping messages. " + "Enforced by app validator and DB CHECK constraint " + f"(pattern {SESSION_ID_SQL_PATTERN}).", + ) # Relationships agent: Mapped["Agent"] = relationship("Agent", back_populates="messages", lazy="selectin") diff --git a/mirix/queue/message.proto b/mirix/queue/message.proto index 9f83ed751..e7c322c98 100644 --- a/mirix/queue/message.proto +++ b/mirix/queue/message.proto @@ -98,6 +98,9 @@ message MessageCreate { // Optional multi-agent group id optional string group_id = 7; + + // Optional top-level session identifier (see mirix.schemas.message.MessageCreate.session_id). + optional string session_id = 8; } // List of message content parts (for structured content) diff --git a/mirix/queue/message_pb2.py b/mirix/queue/message_pb2.py index 32cdfd9c3..fda4612db 100644 --- a/mirix/queue/message_pb2.py +++ b/mirix/queue/message_pb2.py @@ -26,7 +26,7 @@ from google.protobuf import struct_pb2 as google_dot_protobuf_dot_struct__pb2 -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19mirix/queue/message.proto\x12\x05mirix\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xaf\x05\n\x0cQueueMessage\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12,\n\x0einput_messages\x18\x03 \x03(\x0b\x32\x14.mirix.MessageCreate\x12\x15\n\x08\x63haining\x18\x04 \x01(\x08H\x00\x88\x01\x01\x12\x14\n\x07user_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x14\n\x07verbose\x18\x06 \x01(\x08H\x02\x88\x01\x01\x12,\n\x0b\x66ilter_tags\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x16\n\tuse_cache\x18\x08 \x01(\x08H\x03\x88\x01\x01\x12\x18\n\x0boccurred_at\x18\t \x01(\tH\x04\x88\x01\x01\x12\x1e\n\x11langfuse_trace_id\x18\n \x01(\tH\x05\x88\x01\x01\x12$\n\x17langfuse_observation_id\x18\x0b \x01(\tH\x06\x88\x01\x01\x12 \n\x13langfuse_session_id\x18\x0c \x01(\tH\x07\x88\x01\x01\x12\x1d\n\x10langfuse_user_id\x18\r \x01(\tH\x08\x88\x01\x01\x12\x32\n\x11\x62lock_filter_tags\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\x12*\n\x1d\x62lock_filter_tags_update_mode\x18\x0f \x01(\tH\t\x88\x01\x01\x42\x0b\n\t_chainingB\n\n\x08_user_idB\n\n\x08_verboseB\x0c\n\n_use_cacheB\x0e\n\x0c_occurred_atB\x14\n\x12_langfuse_trace_idB\x1a\n\x18_langfuse_observation_idB\x16\n\x14_langfuse_session_idB\x13\n\x11_langfuse_user_idB \n\x1e_block_filter_tags_update_mode\"\xcf\x01\n\x04User\x12\n\n\x02id\x18\x01 \x01(\t\x12\x17\n\x0forganization_id\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x0e\n\x06status\x18\x04 \x01(\t\x12\x10\n\x08timezone\x18\x05 \x01(\t\x12.\n\ncreated_at\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\nis_deleted\x18\x08 \x01(\x08\"\xd9\x02\n\rMessageCreate\x12\'\n\x04role\x18\x01 \x01(\x0e\x32\x19.mirix.MessageCreate.Role\x12\x16\n\x0ctext_content\x18\x02 \x01(\tH\x00\x12\x37\n\x12structured_content\x18\x03 \x01(\x0b\x32\x19.mirix.MessageContentListH\x00\x12\x11\n\x04name\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04otid\x18\x05 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tsender_id\x18\x06 \x01(\tH\x03\x88\x01\x01\x12\x15\n\x08group_id\x18\x07 \x01(\tH\x04\x88\x01\x01\"<\n\x04Role\x12\x14\n\x10ROLE_UNSPECIFIED\x10\x00\x12\r\n\tROLE_USER\x10\x01\x12\x0f\n\x0bROLE_SYSTEM\x10\x02\x42\x0e\n\x0c\x63ontent_typeB\x07\n\x05_nameB\x07\n\x05_otidB\x0c\n\n_sender_idB\x0b\n\t_group_id\">\n\x12MessageContentList\x12(\n\x05parts\x18\x01 \x03(\x0b\x32\x19.mirix.MessageContentPart\"\xbc\x01\n\x12MessageContentPart\x12\"\n\x04text\x18\x01 \x01(\x0b\x32\x12.mirix.TextContentH\x00\x12$\n\x05image\x18\x02 \x01(\x0b\x32\x13.mirix.ImageContentH\x00\x12\"\n\x04\x66ile\x18\x03 \x01(\x0b\x32\x12.mirix.FileContentH\x00\x12-\n\ncloud_file\x18\x04 \x01(\x0b\x32\x17.mirix.CloudFileContentH\x00\x42\t\n\x07\x63ontent\"\x1b\n\x0bTextContent\x12\x0c\n\x04text\x18\x01 \x01(\t\"@\n\x0cImageContent\x12\x10\n\x08image_id\x18\x01 \x01(\t\x12\x13\n\x06\x64\x65tail\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_detail\"\x1e\n\x0b\x46ileContent\x12\x0f\n\x07\x66ile_id\x18\x01 \x01(\t\"*\n\x10\x43loudFileContent\x12\x16\n\x0e\x63loud_file_uri\x18\x01 \x01(\tb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x19mirix/queue/message.proto\x12\x05mirix\x1a\x1fgoogle/protobuf/timestamp.proto\x1a\x1cgoogle/protobuf/struct.proto\"\xaf\x05\n\x0cQueueMessage\x12\x11\n\tclient_id\x18\x01 \x01(\t\x12\x10\n\x08\x61gent_id\x18\x02 \x01(\t\x12,\n\x0einput_messages\x18\x03 \x03(\x0b\x32\x14.mirix.MessageCreate\x12\x15\n\x08\x63haining\x18\x04 \x01(\x08H\x00\x88\x01\x01\x12\x14\n\x07user_id\x18\x05 \x01(\tH\x01\x88\x01\x01\x12\x14\n\x07verbose\x18\x06 \x01(\x08H\x02\x88\x01\x01\x12,\n\x0b\x66ilter_tags\x18\x07 \x01(\x0b\x32\x17.google.protobuf.Struct\x12\x16\n\tuse_cache\x18\x08 \x01(\x08H\x03\x88\x01\x01\x12\x18\n\x0boccurred_at\x18\t \x01(\tH\x04\x88\x01\x01\x12\x1e\n\x11langfuse_trace_id\x18\n \x01(\tH\x05\x88\x01\x01\x12$\n\x17langfuse_observation_id\x18\x0b \x01(\tH\x06\x88\x01\x01\x12 \n\x13langfuse_session_id\x18\x0c \x01(\tH\x07\x88\x01\x01\x12\x1d\n\x10langfuse_user_id\x18\r \x01(\tH\x08\x88\x01\x01\x12\x32\n\x11\x62lock_filter_tags\x18\x0e \x01(\x0b\x32\x17.google.protobuf.Struct\x12*\n\x1d\x62lock_filter_tags_update_mode\x18\x0f \x01(\tH\t\x88\x01\x01\x42\x0b\n\t_chainingB\n\n\x08_user_idB\n\n\x08_verboseB\x0c\n\n_use_cacheB\x0e\n\x0c_occurred_atB\x14\n\x12_langfuse_trace_idB\x1a\n\x18_langfuse_observation_idB\x16\n\x14_langfuse_session_idB\x13\n\x11_langfuse_user_idB \n\x1e_block_filter_tags_update_mode\"\xcf\x01\n\x04User\x12\n\n\x02id\x18\x01 \x01(\t\x12\x17\n\x0forganization_id\x18\x02 \x01(\t\x12\x0c\n\x04name\x18\x03 \x01(\t\x12\x0e\n\x06status\x18\x04 \x01(\t\x12\x10\n\x08timezone\x18\x05 \x01(\t\x12.\n\ncreated_at\x18\x06 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12.\n\nupdated_at\x18\x07 \x01(\x0b\x32\x1a.google.protobuf.Timestamp\x12\x12\n\nis_deleted\x18\x08 \x01(\x08\"\x81\x03\n\rMessageCreate\x12\'\n\x04role\x18\x01 \x01(\x0e\x32\x19.mirix.MessageCreate.Role\x12\x16\n\x0ctext_content\x18\x02 \x01(\tH\x00\x12\x37\n\x12structured_content\x18\x03 \x01(\x0b\x32\x19.mirix.MessageContentListH\x00\x12\x11\n\x04name\x18\x04 \x01(\tH\x01\x88\x01\x01\x12\x11\n\x04otid\x18\x05 \x01(\tH\x02\x88\x01\x01\x12\x16\n\tsender_id\x18\x06 \x01(\tH\x03\x88\x01\x01\x12\x15\n\x08group_id\x18\x07 \x01(\tH\x04\x88\x01\x01\x12\x17\n\nsession_id\x18\x08 \x01(\tH\x05\x88\x01\x01\"<\n\x04Role\x12\x14\n\x10ROLE_UNSPECIFIED\x10\x00\x12\r\n\tROLE_USER\x10\x01\x12\x0f\n\x0bROLE_SYSTEM\x10\x02\x42\x0e\n\x0c\x63ontent_typeB\x07\n\x05_nameB\x07\n\x05_otidB\x0c\n\n_sender_idB\x0b\n\t_group_idB\r\n\x0b_session_id\">\n\x12MessageContentList\x12(\n\x05parts\x18\x01 \x03(\x0b\x32\x19.mirix.MessageContentPart\"\xbc\x01\n\x12MessageContentPart\x12\"\n\x04text\x18\x01 \x01(\x0b\x32\x12.mirix.TextContentH\x00\x12$\n\x05image\x18\x02 \x01(\x0b\x32\x13.mirix.ImageContentH\x00\x12\"\n\x04\x66ile\x18\x03 \x01(\x0b\x32\x12.mirix.FileContentH\x00\x12-\n\ncloud_file\x18\x04 \x01(\x0b\x32\x17.mirix.CloudFileContentH\x00\x42\t\n\x07\x63ontent\"\x1b\n\x0bTextContent\x12\x0c\n\x04text\x18\x01 \x01(\t\"@\n\x0cImageContent\x12\x10\n\x08image_id\x18\x01 \x01(\t\x12\x13\n\x06\x64\x65tail\x18\x02 \x01(\tH\x00\x88\x01\x01\x42\t\n\x07_detail\"\x1e\n\x0b\x46ileContent\x12\x0f\n\x07\x66ile_id\x18\x01 \x01(\t\"*\n\x10\x43loudFileContent\x12\x16\n\x0e\x63loud_file_uri\x18\x01 \x01(\tb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -38,19 +38,19 @@ _globals['_USER']._serialized_start=790 _globals['_USER']._serialized_end=997 _globals['_MESSAGECREATE']._serialized_start=1000 - _globals['_MESSAGECREATE']._serialized_end=1345 - _globals['_MESSAGECREATE_ROLE']._serialized_start=1224 - _globals['_MESSAGECREATE_ROLE']._serialized_end=1284 - _globals['_MESSAGECONTENTLIST']._serialized_start=1347 - _globals['_MESSAGECONTENTLIST']._serialized_end=1409 - _globals['_MESSAGECONTENTPART']._serialized_start=1412 - _globals['_MESSAGECONTENTPART']._serialized_end=1600 - _globals['_TEXTCONTENT']._serialized_start=1602 - _globals['_TEXTCONTENT']._serialized_end=1629 - _globals['_IMAGECONTENT']._serialized_start=1631 - _globals['_IMAGECONTENT']._serialized_end=1695 - _globals['_FILECONTENT']._serialized_start=1697 - _globals['_FILECONTENT']._serialized_end=1727 - _globals['_CLOUDFILECONTENT']._serialized_start=1729 - _globals['_CLOUDFILECONTENT']._serialized_end=1771 + _globals['_MESSAGECREATE']._serialized_end=1385 + _globals['_MESSAGECREATE_ROLE']._serialized_start=1249 + _globals['_MESSAGECREATE_ROLE']._serialized_end=1309 + _globals['_MESSAGECONTENTLIST']._serialized_start=1387 + _globals['_MESSAGECONTENTLIST']._serialized_end=1449 + _globals['_MESSAGECONTENTPART']._serialized_start=1452 + _globals['_MESSAGECONTENTPART']._serialized_end=1640 + _globals['_TEXTCONTENT']._serialized_start=1642 + _globals['_TEXTCONTENT']._serialized_end=1669 + _globals['_IMAGECONTENT']._serialized_start=1671 + _globals['_IMAGECONTENT']._serialized_end=1735 + _globals['_FILECONTENT']._serialized_start=1737 + _globals['_FILECONTENT']._serialized_end=1767 + _globals['_CLOUDFILECONTENT']._serialized_start=1769 + _globals['_CLOUDFILECONTENT']._serialized_end=1811 # @@protoc_insertion_point(module_scope) diff --git a/mirix/queue/message_pb2.pyi b/mirix/queue/message_pb2.pyi index 4822f24f9..1d36fcc53 100644 --- a/mirix/queue/message_pb2.pyi +++ b/mirix/queue/message_pb2.pyi @@ -63,7 +63,7 @@ class User(_message.Message): def __init__(self, id: _Optional[str] = ..., organization_id: _Optional[str] = ..., name: _Optional[str] = ..., status: _Optional[str] = ..., timezone: _Optional[str] = ..., created_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., updated_at: _Optional[_Union[_timestamp_pb2.Timestamp, _Mapping]] = ..., is_deleted: bool = ...) -> None: ... class MessageCreate(_message.Message): - __slots__ = ("role", "text_content", "structured_content", "name", "otid", "sender_id", "group_id") + __slots__ = ("role", "text_content", "structured_content", "name", "otid", "sender_id", "group_id", "session_id") class Role(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): __slots__ = () ROLE_UNSPECIFIED: _ClassVar[MessageCreate.Role] @@ -79,6 +79,7 @@ class MessageCreate(_message.Message): OTID_FIELD_NUMBER: _ClassVar[int] SENDER_ID_FIELD_NUMBER: _ClassVar[int] GROUP_ID_FIELD_NUMBER: _ClassVar[int] + SESSION_ID_FIELD_NUMBER: _ClassVar[int] role: MessageCreate.Role text_content: str structured_content: MessageContentList @@ -86,7 +87,8 @@ class MessageCreate(_message.Message): otid: str sender_id: str group_id: str - def __init__(self, role: _Optional[_Union[MessageCreate.Role, str]] = ..., text_content: _Optional[str] = ..., structured_content: _Optional[_Union[MessageContentList, _Mapping]] = ..., name: _Optional[str] = ..., otid: _Optional[str] = ..., sender_id: _Optional[str] = ..., group_id: _Optional[str] = ...) -> None: ... + session_id: str + def __init__(self, role: _Optional[_Union[MessageCreate.Role, str]] = ..., text_content: _Optional[str] = ..., structured_content: _Optional[_Union[MessageContentList, _Mapping]] = ..., name: _Optional[str] = ..., otid: _Optional[str] = ..., sender_id: _Optional[str] = ..., group_id: _Optional[str] = ..., session_id: _Optional[str] = ...) -> None: ... class MessageContentList(_message.Message): __slots__ = ("parts",) diff --git a/mirix/queue/queue_util.py b/mirix/queue/queue_util.py index deea94891..41b0cac9c 100644 --- a/mirix/queue/queue_util.py +++ b/mirix/queue/queue_util.py @@ -145,6 +145,8 @@ async def put_messages( proto_msg.sender_id = msg.sender_id if msg.group_id: proto_msg.group_id = msg.group_id + if msg.session_id: + proto_msg.session_id = msg.session_id proto_input_messages.append(proto_msg) diff --git a/mirix/queue/worker.py b/mirix/queue/worker.py index 2cde1df6a..f590d6742 100644 --- a/mirix/queue/worker.py +++ b/mirix/queue/worker.py @@ -98,6 +98,7 @@ def _convert_proto_message_to_pydantic(self, proto_msg) -> "MessageCreate": otid=proto_msg.otid if proto_msg.HasField("otid") else None, sender_id=proto_msg.sender_id if proto_msg.HasField("sender_id") else None, group_id=proto_msg.group_id if proto_msg.HasField("group_id") else None, + session_id=proto_msg.session_id if proto_msg.HasField("session_id") else None, filter_tags=None, ) diff --git a/mirix/schemas/message.py b/mirix/schemas/message.py index d3ece01cc..c466b2e21 100644 --- a/mirix/schemas/message.py +++ b/mirix/schemas/message.py @@ -45,6 +45,44 @@ from mirix.system import unpack_message +import re + +# Single source of truth for the session_id column. The DB CHECK constraint in +# mirix/orm/message.py and the SQL migration in scripts/migrate_add_message_session_id.sql +# are derived from these — keep them in sync via tests/test_session_id.py. +SESSION_ID_MAX_LEN = 64 +SESSION_ID_ALLOWED_CHARS = "A-Za-z0-9_-" +# Python-side Pattern: at least 1 char. The length cap is enforced separately so +# we can return a precise error message. +SESSION_ID_PATTERN = f"^[{SESSION_ID_ALLOWED_CHARS}]+$" +# DB-side pattern includes the length cap inline because Postgres regex has no +# "length=max" notion beyond a bounded quantifier. +SESSION_ID_SQL_PATTERN = f"^[{SESSION_ID_ALLOWED_CHARS}]{{1,{SESSION_ID_MAX_LEN}}}$" +_SESSION_ID_RE = re.compile(SESSION_ID_PATTERN) + + +def _validate_session_id(value: Optional[str]) -> Optional[str]: + """Shared validator for the top-level session_id field. + + Accepts None or a non-empty string of [A-Za-z0-9_-], up to SESSION_ID_MAX_LEN. + Rejects empty strings so the DB layer never stores an ambiguous "". + """ + if value is None: + return None + if not isinstance(value, str) or value == "": + raise ValueError("session_id must be a non-empty string or null") + if len(value) > SESSION_ID_MAX_LEN: + raise ValueError( + f"session_id exceeds max length of {SESSION_ID_MAX_LEN}" + ) + if not _SESSION_ID_RE.match(value): + raise ValueError( + f"session_id must match [{SESSION_ID_ALLOWED_CHARS}]+ " + "(letters, digits, '_' or '-')" + ) + return value + + class BaseMessage(OrmMetadataBase): __id_prefix__ = "message" @@ -69,10 +107,23 @@ class MessageCreate(BaseModel): description="The id of the sender of the message, can be an identity id or agent id", ) group_id: Optional[str] = Field(None, description="The multi-agent group that the message was sent in") + session_id: Optional[str] = Field( + None, + description=( + "Optional top-level conversation/session identifier. Messages sharing a session_id " + "belong to the same logical interaction. Must match [a-zA-Z0-9_-]+ and be <= 64 chars." + ), + examples=["sess-xyz"], + ) filter_tags: Optional[Dict[str, Any]] = Field( None, description="Optional tags for filtering and categorizing this message and related memories" ) + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]: data = super().model_dump(**kwargs) if to_orm and "content" in data: @@ -100,6 +151,15 @@ class MessageUpdate(BaseModel): # created_at: Optional[datetime] = Field(None, description="The time the message was created.") tool_calls: Optional[List[OpenAIToolCall,]] = Field(None, description="The list of tool calls requested.") tool_call_id: Optional[str] = Field(None, description="The id of the tool call.") + session_id: Optional[str] = Field( + None, + description="Update the message's top-level session_id. Same validation as MessageCreate.session_id.", + ) + + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) def model_dump(self, to_orm: bool = False, **kwargs) -> Dict[str, Any]: data = super().model_dump(**kwargs) @@ -157,6 +217,13 @@ class Message(BaseMessage): None, description="The id of the sender of the message, can be an identity id or agent id", ) + session_id: Optional[str] = Field( + None, + description=( + "Top-level conversation/session identifier. Messages sharing a session_id belong to " + "the same logical interaction. Indexed column; must match [a-zA-Z0-9_-]+ and be <= 64 chars." + ), + ) # This overrides the optional base orm schema, created_at MUST exist on all messages objects created_at: datetime = Field( default_factory=get_utc_time, @@ -179,6 +246,11 @@ def validate_role(cls, v: str) -> str: assert v in roles, f"Role must be one of {roles}" return v + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + def to_json(self): json_message = vars(self) if json_message["tool_calls"] is not None: @@ -457,6 +529,7 @@ def dict_to_message( name: Optional[str] = None, group_id: Optional[str] = None, tool_returns: Optional[List[ToolReturn]] = None, + session_id: Optional[str] = None, ): """Convert a ChatCompletion message object into a Message object (synced to DB)""" if not created_at: @@ -516,6 +589,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -530,6 +604,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) elif "function_call" in openai_message_dict and openai_message_dict["function_call"] is not None: @@ -565,6 +640,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -579,6 +655,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: @@ -628,6 +705,7 @@ def dict_to_message( id=str(id), tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) else: return Message( @@ -642,6 +720,7 @@ def dict_to_message( created_at=created_at, tool_returns=tool_returns, group_id=group_id, + session_id=session_id, ) def to_openai_dict_search_results(self, max_tool_id_length: int = TOOL_CALL_ID_MAX_LEN) -> dict: diff --git a/mirix/server/rest_api.py b/mirix/server/rest_api.py index a9963df83..51189300c 100644 --- a/mirix/server/rest_api.py +++ b/mirix/server/rest_api.py @@ -17,12 +17,12 @@ from fastapi import APIRouter, Body, FastAPI, Header, HTTPException, Query, Request from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import JSONResponse -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, field_validator, model_validator from mirix.helpers.message_helpers import prepare_input_message_create from mirix.llm_api.llm_client import LLMClient from mirix.log import get_logger -from mirix.orm.errors import NoResultFound +from mirix.orm.errors import NoResultFound, UniqueConstraintViolationError from mirix.schemas.agent import AgentState, AgentType, CreateAgent from mirix.schemas.block import Block, BlockUpdate, CreateBlock, Human, Persona from mirix.schemas.client import Client, ClientCreate, ClientUpdate @@ -36,7 +36,7 @@ from mirix.schemas.file import FileMetadata from mirix.schemas.llm_config import LLMConfig from mirix.schemas.memory import ArchivalMemorySummary, Memory, RecallMemorySummary -from mirix.schemas.message import Message, MessageCreate +from mirix.schemas.message import Message, MessageCreate, _validate_session_id from mirix.schemas.mirix_response import MirixResponse from mirix.schemas.organization import Organization from mirix.schemas.procedural_memory import ProceduralMemoryItemUpdate @@ -1024,6 +1024,7 @@ class SendMessageRequest(BaseModel): role: str user_id: Optional[str] = None # End-user ID for message attribution name: Optional[str] = None + session_id: Optional[str] = None # Top-level session identifier stream_steps: bool = False stream_tokens: bool = False filter_tags: Optional[Dict[str, Any]] = None # Filter tags support @@ -1033,6 +1034,11 @@ class SendMessageRequest(BaseModel): block_filter_tags_update_mode: Optional[str] = "merge" # "merge" or "replace" use_cache: bool = True # Control Redis cache behavior + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + @app.post("/agents/{agent_id}/messages", response_model=MirixResponse) async def send_message_to_agent( @@ -1076,6 +1082,7 @@ async def send_message_to_agent( role=MessageRole(request.role), content=request.message, name=request.name, + session_id=request.session_id, ) # Put message on queue for processing @@ -1994,6 +2001,7 @@ class AddMemoryRequest(BaseModel): chaining: bool = True verbose: bool = False filter_tags: Optional[Dict[str, Any]] = None + session_id: Optional[str] = None # Batch-level session id applied to every message in this request block_filter_tags: Optional[Dict[str, Any]] = ( None # Applied only when blocks are created (e.g. from default template) ) @@ -2001,6 +2009,24 @@ class AddMemoryRequest(BaseModel): use_cache: bool = True # Control Redis cache behavior occurred_at: Optional[str] = None # Optional ISO 8601 timestamp string for episodic memory + @field_validator("session_id") + @classmethod + def _check_session_id(cls, v: Optional[str]) -> Optional[str]: + return _validate_session_id(v) + + @model_validator(mode="after") + def _check_session_id_agrees_with_filter_tags(self) -> "AddMemoryRequest": + # Codex review I1: if both the top-level session_id and filter_tags.session_id are + # provided, they must match. Otherwise message-scope and memory-scope silently diverge. + if self.session_id is not None and isinstance(self.filter_tags, dict): + tag_sid = self.filter_tags.get("session_id") + if tag_sid is not None and tag_sid != self.session_id: + raise ValueError( + "session_id and filter_tags['session_id'] must agree; " + f"got {self.session_id!r} vs {tag_sid!r}" + ) + return self + @router.post("/memory/add") @with_langfuse_tracing @@ -2079,6 +2105,12 @@ async def add_memory( input_messages = convert_message_to_mirix_message(message) + # Batch-level session_id: stamp every message that didn't carry its own. + if request.session_id is not None: + for msg_create in input_messages: + if msg_create.session_id is None: + msg_create.session_id = request.session_id + # Add client scope to filter_tags (create if not provided) if request.filter_tags is not None: # Create a copy to avoid modifying the original request @@ -2087,6 +2119,11 @@ async def add_memory( # Create new filter_tags if not provided filter_tags = {} + # Mirror session_id into filter_tags so extracted memories inherit it too. + # The model_validator on AddMemoryRequest has already ensured agreement if both were set. + if request.session_id is not None: + filter_tags["session_id"] = request.session_id + if request.block_filter_tags is not None and not isinstance(request.block_filter_tags, dict): raise HTTPException(status_code=400, detail="block_filter_tags must be a dict when provided") if request.block_filter_tags is not None: @@ -2184,11 +2221,21 @@ async def add_memory_sync( input_messages = convert_message_to_mirix_message(message) + # Batch-level session_id: stamp every message that didn't carry its own. + if request.session_id is not None: + for msg_create in input_messages: + if msg_create.session_id is None: + msg_create.session_id = request.session_id + if request.filter_tags is not None: filter_tags = dict(request.filter_tags) else: filter_tags = {} + # The AddMemoryRequest model_validator already ensured agreement if both were set. + if request.session_id is not None: + filter_tags["session_id"] = request.session_id + if client.write_scope is None: raise HTTPException(status_code=403, detail="Client has no write_scope - cannot create memories") filter_tags["scope"] = client.write_scope diff --git a/mirix/services/message_manager.py b/mirix/services/message_manager.py index 1ccffdb3d..b0f872dc2 100755 --- a/mirix/services/message_manager.py +++ b/mirix/services/message_manager.py @@ -433,6 +433,7 @@ async def list_user_messages_for_agent( filters: Optional[Dict] = None, query_text: Optional[str] = None, ascending: bool = True, + session_id: Optional[str] = None, ) -> List[PydanticMessage]: """List user messages with flexible filtering and pagination options. @@ -443,6 +444,7 @@ async def list_user_messages_for_agent( limit: Maximum number of records to return filters: Additional filters to apply query_text: Optional text to search for in message content + session_id: Optional session id to filter by (exact match on column). Returns: List[PydanticMessage] - List of messages matching the criteria @@ -461,6 +463,7 @@ async def list_user_messages_for_agent( filters=message_filters, query_text=query_text, ascending=ascending, + session_id=session_id, ) @update_timezone @@ -477,6 +480,7 @@ async def list_messages_for_agent( query_text: Optional[str] = None, ascending: bool = True, use_cache: bool = True, + session_id: Optional[str] = None, ) -> List[PydanticMessage]: """List messages with flexible filtering and pagination options. @@ -501,6 +505,8 @@ async def list_messages_for_agent( message_filters.update({"organization_id": actor.organization_id}) if filters: message_filters.update(filters) + if session_id is not None: + message_filters["session_id"] = session_id results = await MessageModel.list( db_session=session, diff --git a/pyproject.toml b/pyproject.toml index 2477ff926..7a7405313 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -85,7 +85,7 @@ dependencies = [ "ruff (>=0.13.3,<0.14.0)", "pyright (>=1.1.406,<2.0.0)", "ipdb (>=0.13.13,<0.14.0)", - "protobuf (>=5.0.0,<6.0.0)", + "protobuf (>=5.27.2,<6.0.0)", # floor matches gencode header in mirix/queue/message_pb2.py "redis[hiredis] (>=7.0.1,<8.0.0)", # async client only (redis.asyncio); no sync redis "bcrypt (>=4.0.0)", "PyJWT (>=2.10.1)", diff --git a/requirements.txt b/requirements.txt index 90b1f7082..8933f0d4b 100644 --- a/requirements.txt +++ b/requirements.txt @@ -46,7 +46,7 @@ asyncpg aiosqlite httpx ipdb -protobuf>=5.0.0,<6.0.0 +protobuf>=5.27.2,<6.0.0 redis[hiredis]>=7.0.1,<8.0.0 aiokafka>=0.13.0,<0.14.0 asyncddgs diff --git a/scripts/migrate_add_message_session_id.sql b/scripts/migrate_add_message_session_id.sql new file mode 100644 index 000000000..346c5b64a --- /dev/null +++ b/scripts/migrate_add_message_session_id.sql @@ -0,0 +1,54 @@ +-- Migration: add top-level session_id column + index on messages. +-- Run once on existing databases. New databases get this via SQLAlchemy create_all. +-- +-- This migration is split into two phases because CREATE INDEX CONCURRENTLY +-- cannot run inside a transaction block. +-- +-- PHASE 1 (transactional): add the column and a CHECK constraint so any writer +-- is immediately bound by the format rules. Runs quickly; acquires a short +-- ACCESS EXCLUSIVE lock only for the DDL itself. +-- +-- PHASE 2 (must run OUTSIDE a transaction, not inside psql -1): +-- CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_messages_agent_session_created_at +-- ON messages (agent_id, session_id, created_at); +-- +-- CONCURRENTLY avoids the write-blocking ACCESS EXCLUSIVE lock on large tables. +-- If the index build fails midway, Postgres leaves an INVALID index behind; drop +-- it and retry. + +-- Phase 1 is idempotent: safe to re-run (e.g. if phase 2 failed and the whole +-- migration is replayed). Column/constraint additions all guard against +-- existing state. + +BEGIN; + +ALTER TABLE messages + ADD COLUMN IF NOT EXISTS session_id VARCHAR(64); + +-- CHECK constraint matches mirix.schemas.message._validate_session_id. +-- The constraint is NOT VALID first so adding it on a large table does not +-- rewrite existing rows; legacy rows all have NULL session_id and satisfy it. +-- Guarded so a re-run (after a partial failure) is a no-op. +DO $$ +BEGIN + IF NOT EXISTS ( + SELECT 1 FROM pg_constraint + WHERE conname = 'ck_messages_session_id_format' + AND conrelid = 'messages'::regclass + ) THEN + ALTER TABLE messages + ADD CONSTRAINT ck_messages_session_id_format + CHECK (session_id IS NULL OR session_id ~ '^[A-Za-z0-9_-]{1,64}$') + NOT VALID; + END IF; +END$$; + +-- Validate the constraint in a separate step so it only takes SHARE UPDATE +-- EXCLUSIVE (reads + writes continue; no concurrent DDL only). +-- VALIDATE is a no-op if the constraint is already validated. +ALTER TABLE messages VALIDATE CONSTRAINT ck_messages_session_id_format; + +COMMIT; + +-- PHASE 2 — run scripts/migrate_add_message_session_id_phase2.sql separately, +-- NOT inside a transaction (do not combine with this file via `psql -1`). diff --git a/scripts/migrate_add_message_session_id_phase2.sql b/scripts/migrate_add_message_session_id_phase2.sql new file mode 100644 index 000000000..4746ce771 --- /dev/null +++ b/scripts/migrate_add_message_session_id_phase2.sql @@ -0,0 +1,11 @@ +-- Phase 2 of the session_id migration — run AFTER migrate_add_message_session_id.sql. +-- MUST run outside a transaction (no BEGIN/COMMIT here, and do NOT run via `psql -1`). +-- Use e.g.: +-- psql "$DATABASE_URL" -f scripts/migrate_add_message_session_id_phase2.sql +-- +-- CONCURRENTLY avoids the write-blocking ACCESS EXCLUSIVE lock on large tables. +-- If the build fails midway, Postgres leaves an INVALID index behind; drop it and retry: +-- DROP INDEX CONCURRENTLY IF EXISTS ix_messages_agent_session_created_at; + +CREATE INDEX CONCURRENTLY IF NOT EXISTS ix_messages_agent_session_created_at + ON messages (agent_id, session_id, created_at); diff --git a/tests/test_session_id.py b/tests/test_session_id.py new file mode 100644 index 000000000..71c01f1c0 --- /dev/null +++ b/tests/test_session_id.py @@ -0,0 +1,497 @@ +""" +Unit tests for the top-level `session_id` field on messages. + +Scope: +- Pydantic schema: MessageCreate / MessageUpdate / Message accept optional session_id. +- Validation: length/charset rules on session_id. +- Helper: prepare_input_message_create propagates session_id to the Message. +- REST schemas: SendMessageRequest / AddMemoryRequest accept session_id. +- Queue serialization: put_messages path sets session_id on proto, worker restores it. +- ORM column exists and is indexed. + +These are unit tests only; no DB, no running server. +""" + +from __future__ import annotations + +import pytest + +from mirix.helpers.message_helpers import prepare_input_message_create +from mirix.schemas.enums import MessageRole +from mirix.schemas.message import Message, MessageCreate, MessageUpdate + + +# ----------------------------- Schema ------------------------------------ + + +class TestMessageCreateSessionId: + def test_accepts_session_id(self): + m = MessageCreate(role=MessageRole.user, content="hi", session_id="sess-abc") + assert m.session_id == "sess-abc" + + def test_session_id_optional(self): + m = MessageCreate(role=MessageRole.user, content="hi") + assert m.session_id is None + + def test_rejects_empty_string_session_id(self): + # Empty string is ambiguous; require None or a non-empty string. + with pytest.raises(ValueError): + MessageCreate(role=MessageRole.user, content="hi", session_id="") + + def test_rejects_too_long_session_id(self): + with pytest.raises(ValueError): + MessageCreate(role=MessageRole.user, content="hi", session_id="a" * 65) + + def test_rejects_invalid_chars(self): + with pytest.raises(ValueError): + MessageCreate(role=MessageRole.user, content="hi", session_id="sess/abc") + + def test_accepts_allowed_chars(self): + m = MessageCreate( + role=MessageRole.user, + content="hi", + session_id="sess_Abc-123", + ) + assert m.session_id == "sess_Abc-123" + + +class TestMessageSchemaSessionId: + def test_accepts_session_id(self): + m = Message( + agent_id="agent-1", + role=MessageRole.user, + session_id="sess-xyz", + ) + assert m.session_id == "sess-xyz" + + def test_defaults_to_none(self): + m = Message(agent_id="agent-1", role=MessageRole.user) + assert m.session_id is None + + +class TestMessageUpdateSessionId: + def test_accepts_session_id(self): + u = MessageUpdate(session_id="sess-upd") + assert u.session_id == "sess-upd" + + def test_defaults_to_none(self): + u = MessageUpdate() + assert u.session_id is None + + +# ----------------------------- Helper ------------------------------------ + + +class TestPrepareInputMessageCreate: + def test_propagates_session_id(self): + create = MessageCreate(role=MessageRole.user, content="hi", session_id="sess-1") + msg = prepare_input_message_create(create, agent_id="agent-1") + assert msg.session_id == "sess-1" + + def test_missing_session_id_is_none(self): + create = MessageCreate(role=MessageRole.user, content="hi") + msg = prepare_input_message_create(create, agent_id="agent-1") + assert msg.session_id is None + + +class TestDictToMessageSessionId: + """Codex review C1: internal agent-synthesized messages must inherit session_id.""" + + def test_passes_session_id_through_kwarg(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={"role": "user", "content": "hi"}, + session_id="sess-internal", + ) + assert msg.session_id == "sess-internal" + + def test_defaults_to_none(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={"role": "user", "content": "hi"}, + ) + assert msg.session_id is None + + def test_tool_role_message_carries_session_id(self): + msg = Message.dict_to_message( + agent_id="agent-1", + openai_message_dict={ + "role": "tool", + "content": "ok", + "tool_call_id": "tc-1", + "name": "do_stuff", + }, + session_id="sess-internal", + ) + assert msg.session_id == "sess-internal" + + +# ----------------------------- REST request schemas ---------------------- + + +class TestRestRequestSchemasSessionId: + def test_send_message_request_accepts_session_id(self): + from mirix.server.rest_api import SendMessageRequest + + req = SendMessageRequest(message="hi", role="user", session_id="sess-req") + assert req.session_id == "sess-req" + + def test_send_message_request_session_id_optional(self): + from mirix.server.rest_api import SendMessageRequest + + req = SendMessageRequest(message="hi", role="user") + assert req.session_id is None + + def test_add_memory_request_accepts_session_id(self): + from mirix.server.rest_api import AddMemoryRequest + + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-mem", + ) + assert req.session_id == "sess-mem" + + def test_add_memory_request_session_id_optional(self): + from mirix.server.rest_api import AddMemoryRequest + + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + ) + assert req.session_id is None + + # Validation at the REST boundary (Codex review: C2 + I5). + # Invalid input should raise at model construction, not be queued. + def test_send_message_request_rejects_invalid_session_id(self): + from mirix.server.rest_api import SendMessageRequest + + with pytest.raises(ValueError): + SendMessageRequest(message="hi", role="user", session_id="bad/chars") + with pytest.raises(ValueError): + SendMessageRequest(message="hi", role="user", session_id="") + with pytest.raises(ValueError): + SendMessageRequest(message="hi", role="user", session_id="a" * 65) + + def test_add_memory_request_rejects_invalid_session_id(self): + from mirix.server.rest_api import AddMemoryRequest + + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="bad chars", + ) + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="", + ) + + +# ----------------------------- Queue proto ------------------------------- + + +class TestAddMemoryMismatchRejection: + """Codex review I1: top-level session_id and filter_tags.session_id must agree.""" + + def test_matching_values_are_allowed(self): + from mirix.server.rest_api import AddMemoryRequest + + # Same value in both places is fine. + req = AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-1", + filter_tags={"session_id": "sess-1"}, + ) + assert req.session_id == "sess-1" + assert req.filter_tags["session_id"] == "sess-1" + + def test_mismatch_is_rejected_at_request_model(self): + from mirix.server.rest_api import AddMemoryRequest + + with pytest.raises(ValueError): + AddMemoryRequest( + meta_agent_id="meta-1", + messages=[{"role": "user", "content": "hi"}], + session_id="sess-1", + filter_tags={"session_id": "sess-2"}, + ) + + +class TestQueueProtoSessionId: + """Ensure the proto round-trips session_id through MessageCreate.""" + + def test_proto_message_create_has_session_id_field(self): + from mirix.queue.message_pb2 import MessageCreate as ProtoMessageCreate + + field_names = {f.name for f in ProtoMessageCreate.DESCRIPTOR.fields} + assert "session_id" in field_names + + def test_put_messages_serializes_session_id(self, monkeypatch): + """put_messages should copy MessageCreate.session_id onto proto.""" + import asyncio + + from mirix.queue import queue_util + from mirix.schemas.client import Client + + saved = {} + + class FakeQueue: + async def save(self, msg): + saved["msg"] = msg + + monkeypatch.setattr(queue_util, "queue", FakeQueue()) + + async def run(): + await queue_util.put_messages( + actor=Client( + id="client-1", + organization_id="org-1", + name="c", + write_scope="w", + read_scopes=["w"], + ), + agent_id="agent-1", + input_messages=[ + MessageCreate( + role=MessageRole.user, + content="hi", + session_id="sess-q", + ) + ], + ) + + asyncio.run(run()) + + msg = saved["msg"] + assert len(msg.input_messages) == 1 + proto = msg.input_messages[0] + assert proto.HasField("session_id") + assert proto.session_id == "sess-q" + + def test_worker_restores_session_id_from_proto(self): + from mirix.queue.message_pb2 import MessageCreate as ProtoMessageCreate + from mirix.queue.worker import QueueWorker + + proto = ProtoMessageCreate() + proto.role = ProtoMessageCreate.ROLE_USER + proto.text_content = "hi" + proto.session_id = "sess-w" + + worker = QueueWorker.__new__(QueueWorker) + out = worker._convert_proto_message_to_pydantic(proto) + assert out.session_id == "sess-w" + + +# ----------------------------- ORM -------------------------------------- + + +class TestAgentStepPropagation: + """Codex v2 review: heartbeat / warning / meta / summary messages inside Agent.step + must inherit session_id from the triggering input so session-scoped list queries + return the whole conversation, not just the initial user turn.""" + + def test_step_heartbeat_sites_include_session_id(self): + """Sanity check: each heartbeat dict_to_message site in Agent.step passes a session_id kwarg. + + We scan mirix/agent/agent.py for every `Message.dict_to_message(` opener and + require a `session_id=` kwarg before the matching outer `)`. A regex with + lazy-matched parens trips on nested calls (e.g. `get_contine_chaining(...)` + spread across lines after ruff reformat), so we walk parens explicitly. + """ + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + opener = "Message.dict_to_message(" + i = 0 + blocks = [] + while True: + j = src.find(opener, i) + if j == -1: + break + depth = 0 + k = j + len(opener) - 1 # position of the opening `(` + for k in range(j + len(opener) - 1, len(src)): + ch = src[k] + if ch == "(": + depth += 1 + elif ch == ")": + depth -= 1 + if depth == 0: + break + blocks.append(src[j : k + 1]) + i = k + 1 + + assert blocks, "expected at least one Message.dict_to_message call in agent.py" + missing = [b for b in blocks if "session_id=" not in b] + assert not missing, ( + "These dict_to_message sites do not pass session_id:\n" + + "\n---\n".join(missing) + ) + + def test_persisted_message_create_sites_include_session_id(self): + """Codex v3 follow-up: MessageCreate(...) sites that are persisted (turned + into a Message via prepare_input_message_create) must also inherit session_id. + + agent.py has two relevant MessageCreate() call sites: the meta-memory + bootstrap (persisted via prepare_input_message_create) and topic-extraction + scratch prompts (sent straight to the LLM, not persisted). We lint the + meta-bootstrap path specifically. + """ + import re + from pathlib import Path + + src = Path("mirix/agent/agent.py").read_text() + # The meta-bootstrap MessageCreate is inside the `if ... meta_memory_agent ...` branch. + meta_block_match = re.search( + r"meta_message\s*=\s*prepare_input_message_create\(\n" + r"\s*MessageCreate\([\s\S]*?\)\s*,\n", + src, + ) + assert meta_block_match, "meta-bootstrap MessageCreate block not found" + assert "session_id=step_session_id" in meta_block_match.group(0), ( + "meta-memory bootstrap MessageCreate must carry session_id=step_session_id, " + f"got:\n{meta_block_match.group(0)}" + ) + + def test_step_seeds_and_restores_stash(self): + """Codex v4 review: Agent.step() also stashes _current_step_session_id on + self for the duration of the step. Without a try/finally, an exception or a + successful return would leave the prior caller's value clobbered, so a + later summarize_messages_inplace() outside any step() would read a stale + session. Lock in the save/restore-with-finally pattern via source inspection + (a real step() requires LLM + DB so behavioral test is impractical here).""" + import inspect + + from mirix.agent import agent as agent_mod + + src = inspect.getsource(agent_mod.Agent.step) + assert ( + 'prev_session_id = getattr(self, "_current_step_session_id", None)' in src + ), ( + "step() must capture the prior _current_step_session_id before overwriting it" + ) + assert "self._current_step_session_id = step_session_id" in src + assert "try:" in src and "finally:" in src, "step() must use try/finally" + assert "self._current_step_session_id = prev_session_id" in src, ( + "step() must restore prev_session_id in its finally block" + ) + + +class TestOrmSessionId: + def test_orm_has_session_id_column(self): + from mirix.orm.message import Message as MessageORM + + cols = {c.name for c in MessageORM.__table__.columns} + assert "session_id" in cols + + def test_orm_session_id_is_indexed(self): + from mirix.orm.message import Message as MessageORM + + # At least one composite or single index covers session_id. + covered = any( + any(c.name == "session_id" for c in idx.columns) + for idx in MessageORM.__table__.indexes + ) + assert covered, "session_id should be indexed" + + +class TestStepUserMessageSessionContext: + """Codex review v3: step_user_message() bypasses step() and must still seed + _current_step_session_id for any pre-persist summary, then restore the prior value. + We inspect the source rather than run the coroutine (full step needs LLM + DB).""" + + def test_step_user_message_seeds_and_restores_stash(self): + import inspect + + from mirix.agent import agent as agent_mod + + src = inspect.getsource(agent_mod.Agent.step_user_message) + # Seeds current step session context for summarizer. + assert "self._current_step_session_id = session_id" in src + # Uses try/finally to restore the prior value — no leaks across calls. + assert "prev_session_id" in src + assert "finally" in src + assert "self._current_step_session_id = prev_session_id" in src + + def test_summarize_messages_inplace_accepts_explicit_session_id(self): + import inspect + + from mirix.agent import agent as agent_mod + + sig = inspect.signature(agent_mod.Agent.summarize_messages_inplace) + assert "session_id" in sig.parameters + # Default is optional None so existing callers don't break. + assert sig.parameters["session_id"].default is None + + +class TestCheckConstraintDialectGating: + """The CHECK uses PG's `~` regex operator, which SQLite doesn't understand. + Verify the constraint is only emitted for Postgres so SQLite create_all works.""" + + def test_check_compiles_for_postgresql(self): + from sqlalchemy.dialects import postgresql + from sqlalchemy.schema import CreateTable + + from mirix.orm.message import Message as MessageORM + + ddl = str( + CreateTable(MessageORM.__table__).compile(dialect=postgresql.dialect()) + ) + assert "ck_messages_session_id_format" in ddl + assert "~" in ddl # Postgres regex operator + + def test_check_is_suppressed_for_sqlite(self): + from sqlalchemy.dialects import sqlite + from sqlalchemy.schema import CreateTable + + from mirix.orm.message import Message as MessageORM + + ddl = str(CreateTable(MessageORM.__table__).compile(dialect=sqlite.dialect())) + # Constraint must not appear on SQLite — its `~` regex would break CREATE TABLE. + assert "ck_messages_session_id_format" not in ddl + assert "session_id ~" not in ddl + + +class TestSessionIdConstantsInSync: + """Codex v2 nit: one source of truth for pattern/length across Python, ORM CheckConstraint, and SQL.""" + + def test_orm_column_length_matches_constant(self): + from mirix.orm.message import Message as MessageORM + from mirix.schemas.message import SESSION_ID_MAX_LEN + + col = MessageORM.__table__.c.session_id + assert col.type.length == SESSION_ID_MAX_LEN + + def test_orm_check_constraint_uses_shared_pattern(self): + from mirix.orm.message import Message as MessageORM + from mirix.schemas.message import SESSION_ID_SQL_PATTERN + + ck_texts = [ + str(c.sqltext) + for c in MessageORM.__table__.constraints + if getattr(c, "name", None) == "ck_messages_session_id_format" + ] + assert ck_texts, "ck_messages_session_id_format not found" + assert SESSION_ID_SQL_PATTERN in ck_texts[0] + + def test_migration_sql_uses_shared_pattern_and_length(self): + from pathlib import Path + + from mirix.schemas.message import ( + SESSION_ID_MAX_LEN, + SESSION_ID_SQL_PATTERN, + ) + + phase1 = Path("scripts/migrate_add_message_session_id.sql").read_text() + phase2 = Path("scripts/migrate_add_message_session_id_phase2.sql").read_text() + + # Column length must match. + assert f"VARCHAR({SESSION_ID_MAX_LEN})" in phase1 + # CHECK regex must match the shared SQL pattern exactly. + assert SESSION_ID_SQL_PATTERN in phase1 + # Phase 2 must use CONCURRENTLY (online, write-safe). + assert "CREATE INDEX CONCURRENTLY" in phase2 diff --git a/tests/test_session_id_integration.py b/tests/test_session_id_integration.py new file mode 100644 index 000000000..07eebc9d8 --- /dev/null +++ b/tests/test_session_id_integration.py @@ -0,0 +1,213 @@ +""" +Integration test for the top-level session_id field: end-to-end round-trip +through MessageManager against real Postgres. + +Verifies: +- A message created with session_id persists that column. +- list_messages_for_agent(session_id=X) returns only messages for session X. +- list_messages_for_agent(session_id=None) returns all messages for that agent. +- Different session_ids correctly isolate messages. + +Requires the docker-compose Postgres (port 5433). Run: + pytest tests/test_session_id_integration.py -v -m integration +""" +from __future__ import annotations + +import asyncio +import sys +import uuid +from pathlib import Path + +import pytest +import pytest_asyncio + +from mirix.settings import settings + +pytestmark = [ + pytest.mark.integration, + pytest.mark.asyncio(loop_scope="module"), + # Round-trip needs real Postgres; SQLite via aiosqlite throws MissingGreenlet + # in a fresh sync session. Mirrors tests/test_agent_trigger_state_integration.py. + pytest.mark.skipif( + not settings.mirix_pg_uri_no_default, + reason="session_id round-trip needs Postgres (set MIRIX_PG_URI)", + ), +] + +project_root = Path(__file__).parent.parent +sys.path.insert(0, str(project_root)) + + +@pytest_asyncio.fixture(scope="module") +def event_loop(): + loop = asyncio.new_event_loop() + yield loop + loop.close() + + +from mirix.schemas.client import Client as PydanticClient +from mirix.schemas.enums import MessageRole +from mirix.schemas.message import Message as PydanticMessage +from mirix.schemas.mirix_message_content import TextContent +from mirix.schemas.user import User as PydanticUser +from mirix.services.message_manager import MessageManager + + +@pytest.fixture +def message_manager(): + return MessageManager() + + +@pytest_asyncio.fixture(scope="module") +async def test_actor(): + from mirix.schemas.organization import Organization as PydanticOrganization + from mirix.services.client_manager import ClientManager + from mirix.services.organization_manager import OrganizationManager + + org_mgr = OrganizationManager() + client_mgr = ClientManager() + + org_id = f"test-session-id-org-{uuid.uuid4().hex[:8]}" + try: + await org_mgr.get_organization_by_id(org_id) + except Exception: + await org_mgr.create_organization( + PydanticOrganization(id=org_id, name="Session ID Test Org") + ) + + client_id = f"test-session-id-client-{uuid.uuid4().hex[:8]}" + try: + return await client_mgr.get_client_by_id(client_id) + except Exception: + return await client_mgr.create_client( + PydanticClient( + id=client_id, + organization_id=org_id, + name="Session ID Test Client", + write_scope="test-sid", + read_scopes=["test-sid"], + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def test_user(test_actor): + from mirix.services.user_manager import UserManager + + user_mgr = UserManager() + user_id = f"test-session-id-user-{uuid.uuid4().hex[:8]}" + try: + return await user_mgr.get_user_by_id(user_id) + except Exception: + return await user_mgr.create_user( + PydanticUser( + id=user_id, + name="Session ID Test User", + organization_id=test_actor.organization_id, + timezone="UTC", + ) + ) + + +@pytest_asyncio.fixture(scope="module") +async def test_agent(test_actor): + """Create a minimal agent we can attach messages to. + + create_agent rejects null llm_config/embedding_config, so supply + placeholder configs (we never invoke the model — only persist messages). + Pattern mirrors tests/test_redis_integration.py::test_agent. + """ + from mirix.schemas.agent import AgentType, CreateAgent + from mirix.schemas.embedding_config import EmbeddingConfig + from mirix.schemas.llm_config import LLMConfig + from mirix.services.agent_manager import AgentManager + + agent_mgr = AgentManager() + agent = await agent_mgr.create_agent( + agent_create=CreateAgent( + name=f"test-session-agent-{uuid.uuid4().hex[:8]}", + agent_type=AgentType.chat_agent, + description="Test agent for session_id round-trip", + system=None, + llm_config=LLMConfig( + model="gpt-4", + model_endpoint_type="openai", + model_endpoint="https://api.openai.com", + context_window=8192, + ), + embedding_config=EmbeddingConfig( + embedding_model="text-embedding-ada-002", + embedding_endpoint_type="openai", + embedding_dim=1536, + ), + ), + actor=test_actor, + ) + return agent + + +async def _create_msg(mgr, actor, agent, text, session_id): + return await mgr.create_message( + pydantic_msg=PydanticMessage( + agent_id=agent.id, + role=MessageRole.user, + content=[TextContent(text=text)], + session_id=session_id, + ), + actor=actor, + use_cache=False, + ) + + +class TestSessionIdRoundTrip: + async def test_persists_and_filters_by_session_id( + self, message_manager, test_actor, test_agent + ): + sid_a = f"sess-a-{uuid.uuid4().hex[:6]}" + sid_b = f"sess-b-{uuid.uuid4().hex[:6]}" + + a1 = await _create_msg(message_manager, test_actor, test_agent, "A1", sid_a) + a2 = await _create_msg(message_manager, test_actor, test_agent, "A2", sid_a) + b1 = await _create_msg(message_manager, test_actor, test_agent, "B1", sid_b) + null_msg = await _create_msg( + message_manager, test_actor, test_agent, "no-session", None + ) + + # Session A returns only A messages. + got_a = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + session_id=sid_a, + limit=100, + use_cache=False, + ) + ids_a = {m.id for m in got_a} + assert a1.id in ids_a + assert a2.id in ids_a + assert b1.id not in ids_a + assert null_msg.id not in ids_a + + # Session B returns only B. + got_b = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + session_id=sid_b, + limit=100, + use_cache=False, + ) + ids_b = {m.id for m in got_b} + assert b1.id in ids_b + assert a1.id not in ids_b + assert null_msg.id not in ids_b + + # No filter returns everything we created (and the session_id column round-trips). + all_msgs = await message_manager.list_messages_for_agent( + agent_id=test_agent.id, + actor=test_actor, + limit=100, + use_cache=False, + ) + by_id = {m.id: m for m in all_msgs} + assert by_id[a1.id].session_id == sid_a + assert by_id[b1.id].session_id == sid_b + assert by_id[null_msg.id].session_id is None