Skip to content

Add top-level session_id field to messages#126

Open
nanxingw wants to merge 1 commit into
mainfrom
feat/message-session-id
Open

Add top-level session_id field to messages#126
nanxingw wants to merge 1 commit into
mainfrom
feat/message-session-id

Conversation

@nanxingw
Copy link
Copy Markdown
Collaborator

@nanxingw nanxingw commented Apr 28, 2026

Summary

Introduce an optional session_id column on the messages table so callers can group messages into logical conversation sessions independent of agent_id. Validated at three layers (Pydantic, REST request models, Postgres CHECK), propagated end-to-end through the agent's synthesized messages, and exposed through MessageManager.list_messages_for_*(session_id=...) for filtering.

Behavior

  • Optional — existing callers untouched. Unset → NULL. Old rows stay NULL after migration.
  • Strict format — when set, must match [A-Za-z0-9_-]{1,64}. Rejected at Pydantic construction, REST request validation, and Postgres CHECK constraint.
  • FilterableMessageManager.list_messages_for_agent(session_id="...") and list_messages_for_user(session_id="...") return only that session's messages.
  • Auto-mirror to memory tagsadd_memory mirrors the batch-level session_id into filter_tags["session_id"] so memories created from that batch inherit the tag, making them retrievable via the existing /memory/retrieve/conversation endpoint without API surface changes.
  • Agent-internal propagation — every synthesized message in a step (LLM-generated assistant replies, tool returns, heartbeat, warning, meta-memory bootstrap, summary) inherits the triggering input's session_id via a _current_step_session_id stash on self, save/restored with try/finally so it cannot leak across calls.

Test plan

  • 38 unit tests in tests/test_session_id.py (DB-free): schema, validator, helper, REST request schemas, queue proto serialization, ORM column shape and indexing, agent-source lint for propagation sites, dialect-gated CHECK, constants-in-sync.
  • 1 Postgres integration test in tests/test_session_id_integration.py — auto-skips without MIRIX_PG_URI (mirrors tests/test_agent_trigger_state_integration.py pattern).
  • Manual E2E with real OpenRouter LLM: Phase A direct MessageManager round-trip + Phase B validator rejection + Phase C real server.send_messages → _step with LLM-generated tool call. All three synthesized messages (user + assistant + tool result) inherited session_id correctly; negative control verified no cross-session leak.
  • Codex review — first pass surfaced one shippable medium finding (step() stash leak); fixed with try/finally + regression test test_step_seeds_and_restores_stash. Second finding was a pre-existing unrelated bug — see [bug] Agent.step_user_message → inner_step missing required first_input_messge arg #127.

Reviewer hints

  • Use git diff -w on mirix/agent/agent.py — the bulk of that file's diff is body re-indent to wrap step() in try/finally. The actual session_id additions are ~70 lines; the rest is the indent shift.
  • The CHECK constraint in mirix/orm/message.py uses ddl_if(dialect="postgresql") because SQLite has no ~ regex operator. SQLite path verified to suppress the constraint (test_check_is_suppressed_for_sqlite).
  • Shared length / charset constants in mirix/schemas/message.py are the single source of truth — the migration SQL and ORM CHECK both reference them, with test_migration_sql_uses_shared_pattern_and_length keeping them in sync.

Out of scope

🤖 Generated with Claude Code

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) <noreply@anthropic.com>
@CLAassistant
Copy link
Copy Markdown

CLA assistant check
Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
You have signed the CLA already but the status is still pending? Let us recheck it.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants