Escape reserved tokens in untrusted chat inputs

chat_template.jinja

@@ -1,7 +1,19 @@
+{%- macro sanitize_input(value) -%}
+  {{- value | replace('<|', '&lt;|') | replace('|>', '|&gt;') -}}
+{%- endmacro -%}
+
+{%- macro render_text(msg, text) -%}
+  {%- if msg['role'] == 'user' or msg['role'] == 'tool' -%}
+    {{- sanitize_input(text) -}}
+  {%- else -%}
+    {{- text -}}
+  {%- endif -%}
+{%- endmacro -%}
+
 {%- macro render_content(msg) -%}
     {%- set c = msg.get('content') -%}
     {%- if c is string -%}
-      {{ c }}
+      {{ render_text(msg, c) }}
     {%- elif c is not none -%}
       {% for content in c -%}
         {% if content['type'] == 'image' or content['type'] == 'image_url' -%}
@@ -9,7 +21,7 @@
         {% elif content['type'] == 'video' or content['type']== 'video_url'-%}
           <|kimi_k25_video_placeholder|>
         {% else -%}
-          {{ content['text'] }}
+          {{ render_text(msg, content['text']) }}
         {%- endif -%}
       {%- endfor -%}
     {%- endif -%}
@@ -17,11 +29,11 @@
 {% macro set_roles(message) -%}
   {%- set role_name =  message.get('name') or  message['role'] -%}
   {%- if message['role'] == 'user' -%}
-    <|im_user|>{{role_name}}<|im_middle|>
+    <|im_user|>{{sanitize_input(role_name)}}<|im_middle|>
   {%- elif message['role'] == 'assistant' -%}
     <|im_assistant|>{{role_name}}<|im_middle|>
   {%- else -%}
-    <|im_system|>{{role_name}}<|im_middle|>
+    <|im_system|>{{sanitize_input(role_name)}}<|im_middle|>
   {%- endif -%}
 {%- endmacro -%}
 {%- macro render_toolcalls(message) -%}
@@ -48,7 +60,7 @@
       {{render_toolcalls(message)}}
     {%- endif -%}
   {%- elif message['role'] == 'tool' -%}
-    {%- set tool_call_id = message.tool_call_id -%}
+    {%- set tool_call_id = sanitize_input(message.tool_call_id) -%}
     ## Return of {{ tool_call_id }}
 {{render_content(message)}}
   {%- elif message['content'] is not none -%}