fix(validator): reject any non-user/tool trailing message, not just assistant

[Osamaali313]

Problem

MistralRequestValidator._validate_last_message (serving / test mode) is meant to require the conversation's last message to be a user or tool message, or an assistant message with prefix / continue_final_message set — as stated in its own error message and the comment "The last message must be a user or tool message in serving mode...".

But the guard ANDs bad_role with an assistant-only flag:

bad_assistant = isinstance(message, AssistantMessage) and not message.prefix and not continue_final_message
bad_role = message.role not in {Roles.user, Roles.tool}
if bad_assistant and bad_role:
    raise InvalidMessageStructureException("Expected last role User or Tool ... for serving but got ...")

bad_assistant is only ever True for an AssistantMessage, so the condition can never fire for any other role. Since _validate_message_order permits user → system, a conversation ending in a SystemMessage reaches this check and passes — even though it's exactly what the message says should be rejected.

Reproduced (serving mode):

v = MistralRequestValidator(mode=ValidationMode.serving)
v.validate_messages([UserMessage(content="hi"), SystemMessage(content="sys")], continue_final_message=False)
# current: accepted (no exception)   expected: InvalidMessageStructureException

The test suite reveals the intent: every case that ends a conversation with a SystemMessage appends a trailing UserMessage with the comment "so we don't get an error for ending with a system message" (tests/validation/test_chat_validation.py:82, :101). The maintainers assume a trailing system message is rejected — but the and bug means it isn't. No test asserts a trailing system message is accepted.

Fix

Reject when the role is not user/tool unless the message is a valid trailing assistant (prefix or continue_final_message):

bad_role = message.role not in {Roles.user, Roles.tool}
valid_trailing_assistant = isinstance(message, AssistantMessage) and (message.prefix or continue_final_message)
if bad_role and not valid_trailing_assistant:
    raise ...

Verified against all role × prefix × continue_final_message combinations — only assistant-without-prefix and system-last now raise; user, tool, assistant-with-prefix, and the continue_final_message path are unchanged (and the following elif continue_final_message branch still behaves identically). A naive bad_role or bad_assistant was avoided because it would wrongly reject a valid prefix=True assistant.

Tests

Adds test_ends_with_system asserting a trailing system message raises in serving mode. The full test_chat_validation.py message-structure suite passes (the only failures in my environment are the pre-existing test_audio_* cases that require the optional soundfile extra, unrelated to this change).

[Copilot]

Pull request overview

Note

Copilot was unable to run its full agentic suite in this review.

Fixes validation of the last chat message role in “serving” mode so that invalid trailing roles (e.g., system) are properly rejected, and adds regression coverage for the bug.

Changes:

• Corrects the last-message validation logic to reject non-user/tool roles unless the last message is a continuable assistant.
• Adds a regression test to ensure a trailing system message raises the expected exception.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
tests/validation/test_chat_validation.py	Adds a regression test for trailing system message rejection.
src/mistral_common/protocol/instruct/validator.py	Fixes boolean logic so invalid trailing roles are rejected in serving mode.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Osamaali313]

Thanks @copilot — two good questions; I looked at both:

1. continue_final_message in valid_trailing_assistant. This preserves the original semantics exactly. The old guard was bad_assistant = isinstance(message, AssistantMessage) and not message.prefix and not continue_final_message, i.e. an assistant was valid as the last message when prefix or continue_final_message — which is precisely valid_trailing_assistant. The prefix-vs-continuation interaction is then enforced by the next branch, which is unchanged:

elif continue_final_message and (last_message_role != Roles.assistant or message.prefix):
    raise ...  # "Expected last role Assistant with prefix False ... continue_final_message"

So with continue_final_message=True: a prefix=True assistant is rejected by the elif, and only a prefix=False assistant is accepted. Removing continue_final_message from valid_trailing_assistant would actually break the valid case — the if would then reject a prefix=False assistant that continuation legitimately allows. I verified all six role × prefix × continue combinations; only assistant-without-prefix and system-last raise, matching intent.

2. Strict regex in the new test. I matched the full message on purpose, to stay consistent with the existing test_ends_with_assistant in this file, which asserts the same full string. Happy to switch both to a focused substring (e.g. r"but got system" / r"but got assistant") if you'd prefer the looser style across these tests — just let me know and I'll update them together.