llmcompressor.modeling.kimi_k3.encoding_k3

def normalize_xtml_tool_result_messages(messages: list[Any]) -> list[Any]:
    """Re-sort K3 XTML tool results into assistant ``tool_calls`` order.

    Serving frameworks generally deliver tool results already in call order. A
    direct Transformers caller, however, may pass OpenAI-style tool messages in any
    order, so each run of consecutive tool messages is matched against the most
    recent preceding assistant ``tool_calls`` by opaque ``tool_call_id`` ==
    ``tool_calls[].id`` (K3 drops the ``func:index`` format requirement) and
    sorted by the matched 1-based position. The matched call is authoritative,
    so each matched message's ``tool`` is set to that call's function name --
    this keeps an explicit (and possibly stale) ``tool``/``name`` from drifting
    out of sync with the reordered position. ``index`` is still derived from the
    rendered position by the chat template. A run that cannot be fully matched is
    left untouched. Re-running is idempotent.

    This function is side-effect free: matched tool messages are shallow-copied
    before their ``tool``/``name`` is rewritten, and every other message is
    appended to the output as-is. The input list and its message objects are
    never mutated.
    """
    if not isinstance(messages, list):
        return messages

    output: list[Any] = []
    current_index: dict = {}
    i = 0
    n = len(messages)

    while i < n:
        message = messages[i]

        if isinstance(message, dict) and message.get("role") == "assistant":
            tool_calls = message.get("tool_calls")
            current_index = _tool_call_id_index(tool_calls) if tool_calls else {}
            output.append(message)
            i += 1
            continue

        if not isinstance(message, dict) or message.get("role") != "tool":
            output.append(message)
            i += 1
            continue

        run: list[tuple] = []  # (position, original_offset, message, name)
        unresolved = False
        offset = 0
        while (
            i < n
            and isinstance(messages[i], dict)
            and messages[i].get("role") == "tool"
        ):
            tool_message = messages[i]
            call_id = tool_message.get("tool_call_id", tool_message.get("id"))
            matched = current_index.get(str(call_id)) if call_id is not None else None
            if matched is None:
                unresolved = True
                run.append((None, offset, tool_message, None))
            else:
                position, name = matched
                run.append((position, offset, tool_message, name))
            offset += 1
            i += 1

        if unresolved:
            output.extend(item[2] for item in run)
        else:
            run.sort(key=lambda item: (item[0], item[1]))
            for _, _, tool_message, name in run:
                if name is None:
                    output.append(tool_message)
                    continue
                # The id-matched call is authoritative: align tool (and any
                # explicit name) so the rendered XTML tool attribute cannot
                # disagree with the reordered position. Copy first so the
                # caller's message object is never mutated.
                resolved = dict(tool_message)
                resolved["tool"] = name
                if "name" in resolved:
                    resolved["name"] = name
                output.append(resolved)

    return output