[Frontend] Cohere chat v2 api support (#47189)

Signed-off-by: Andrew Berneshawi <andrewberneshawi@cohere.com>
Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>

requirements/test/cpu.txt

@@ -126,6 +126,8 @@ click==8.4.2
     #   uvicorn
 cloudpickle==3.1.2
     # via -r requirements/test/../common.txt
+cohere==7.0.8
+    # via -r requirements/test/cuda.in
 cohere-melody==0.9.0
     # via -r requirements/test/cuda.in
 colorama==0.4.6
@@ -226,6 +228,8 @@ fastar==0.11.0
     # via
     #   fastapi
     #   fastapi-cloud-cli
+fastavro==1.12.2
+    # via cohere
 fastparquet==2024.11.0
     # via genai-perf
 fastrlock==0.8.2
@@ -326,6 +330,7 @@ httpx==0.27.2
     # via
     #   -r requirements/test/cuda.in
     #   anthropic
+    #   cohere
     #   fastapi
     #   fastapi-cloud-cli
     #   huggingface-hub
@@ -772,6 +777,7 @@ pydantic==2.12.0
     #   -r requirements/test/../common.txt
     #   albumentations
     #   anthropic
+    #   cohere
     #   compressed-tensors
     #   datamodel-code-generator
     #   fastapi
@@ -789,7 +795,9 @@ pydantic==2.12.0
     #   ray
     #   xgrammar
 pydantic-core==2.41.1
-    # via pydantic
+    # via
+    #   cohere
+    #   pydantic
 pydantic-extra-types==2.10.5
     # via
     #   fastapi
@@ -911,6 +919,7 @@ requests==2.32.3
     #   -r requirements/test/../common.txt
     #   azure-core
     #   buildkite-test-collector
+    #   cohere
     #   datasets
     #   docker
     #   evaluate
@@ -1110,6 +1119,7 @@ tokenizers==0.22.2
     # via
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in
+    #   cohere
     #   transformers
 torch==2.13.0+cpu
     # via
@@ -1191,6 +1201,8 @@ typer==0.26.8
     #   fastsafetensors
     #   perceptron
     #   transformers
+types-requests==2.33.0.20260712
+    # via cohere
 typing-extensions==4.15.0
     # via
     #   -r requirements/test/../common.txt
@@ -1204,6 +1216,7 @@ typing-extensions==4.15.0
     #   azure-identity
     #   azure-storage-blob
     #   chz
+    #   cohere
     #   fastapi
     #   grpcio
     #   huggingface-hub
@@ -1248,6 +1261,7 @@ urllib3==2.2.3
     #   responses
     #   sentry-sdk
     #   tritonclient
+    #   types-requests
 uvicorn==0.35.0
     # via
     #   fastapi

requirements/test/cuda.in

@@ -73,6 +73,7 @@ gpt-oss >= 0.0.7; python_version > '3.11'
 perceptron # required for isaac test
 kaldi-native-fbank >= 1.18.7 # required for fireredasr2 test
 cohere_melody>=0.9.0 # required for cohere command reasoning parser test
+cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types)
 
 # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
 # Older versions are in conflict with teerratorch requirements.

requirements/test/cuda.txt

@@ -131,6 +131,8 @@ click==8.4.2
     #   uvicorn
 cloudpickle==3.1.2
     # via -r requirements/test/../common.txt
+cohere==7.0.8
+    # via -r requirements/test/cuda.in
 cohere-melody==0.9.0
     # via -r requirements/test/cuda.in
 colorama==0.4.6
@@ -242,6 +244,8 @@ fastar==0.11.0
     # via
     #   fastapi
     #   fastapi-cloud-cli
+fastavro==1.12.2
+    # via cohere
 fastparquet==2024.11.0
     # via genai-perf
 fastrlock==0.8.2
@@ -345,6 +349,7 @@ httpx==0.27.2
     # via
     #   -r requirements/test/cuda.in
     #   anthropic
+    #   cohere
     #   fastapi
     #   fastapi-cloud-cli
     #   huggingface-hub
@@ -860,6 +865,7 @@ pydantic==2.12.0
     #   -r requirements/test/../common.txt
     #   albumentations
     #   anthropic
+    #   cohere
     #   compressed-tensors
     #   datamodel-code-generator
     #   fastapi
@@ -877,7 +883,9 @@ pydantic==2.12.0
     #   ray
     #   xgrammar
 pydantic-core==2.41.1
-    # via pydantic
+    # via
+    #   cohere
+    #   pydantic
 pydantic-extra-types==2.10.5
     # via
     #   fastapi
@@ -1002,6 +1010,7 @@ requests==2.32.3
     #   -r requirements/test/../common.txt
     #   azure-core
     #   buildkite-test-collector
+    #   cohere
     #   datasets
     #   docker
     #   evaluate
@@ -1207,6 +1216,7 @@ tokenizers==0.22.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/cuda.in
+    #   cohere
     #   transformers
 torch==2.13.0+cu130
     # via
@@ -1296,6 +1306,8 @@ typer==0.26.8
     #   fastsafetensors
     #   perceptron
     #   transformers
+types-requests==2.33.0.20260712
+    # via cohere
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt
@@ -1310,6 +1322,7 @@ typing-extensions==4.15.0
     #   azure-identity
     #   azure-storage-blob
     #   chz
+    #   cohere
     #   fastapi
     #   grpcio
     #   huggingface-hub
@@ -1354,6 +1367,7 @@ urllib3==2.2.3
     #   responses
     #   sentry-sdk
     #   tritonclient
+    #   types-requests
 uvicorn==0.35.0
     # via
     #   fastapi

requirements/test/rocm.in

@@ -71,6 +71,7 @@ gpt-oss>=0.0.7; python_version > '3.11'
 perceptron # required for isaac test
 kaldi-native-fbank>=1.18.7 # required for fireredasr2 test
 cohere_melody>=0.9.0 # required for cohere command reasoning parser test
+cohere>=7.0.0 # required for cohere chat v2 api tests (protocol/serving import cohere.types)
 
 # Newer versions of datasets require torchcoded, that makes the tests fail in CI because of a missing library.
 # Older versions are in conflict with terratorch requirements.

requirements/test/rocm.txt

@@ -131,6 +131,8 @@ cloudpickle==3.1.2
     # via
     #   -r requirements/test/../common.txt
     #   tilelang
+cohere==7.0.8
+    # via -r requirements/test/rocm.in
 cohere-melody==0.9.0
     # via -r requirements/test/rocm.in
 colorama==0.4.6
@@ -231,6 +233,8 @@ fastapi-cloud-cli==0.16.1
     # via fastapi-cli
 fastar==0.10.0
     # via fastapi-cloud-cli
+fastavro==1.12.2
+    # via cohere
 fastparquet==2026.3.0
     # via genai-perf
 fastsafetensors==0.3.2
@@ -336,6 +340,7 @@ httpx==0.27.2
     # via
     #   -r requirements/test/rocm.in
     #   anthropic
+    #   cohere
     #   fastapi
     #   fastapi-cloud-cli
     #   huggingface-hub
@@ -832,6 +837,7 @@ pydantic==2.12.5
     #   -r requirements/test/../common.txt
     #   albumentations
     #   anthropic
+    #   cohere
     #   compressed-tensors
     #   datamodel-code-generator
     #   fastapi
@@ -849,7 +855,9 @@ pydantic==2.12.5
     #   ray
     #   xgrammar
 pydantic-core==2.41.5
-    # via pydantic
+    # via
+    #   cohere
+    #   pydantic
 pydantic-extra-types==2.11.1
     # via
     #   fastapi
@@ -978,6 +986,7 @@ requests==2.32.5
     #   -r requirements/test/../common.txt
     #   azure-core
     #   buildkite-test-collector
+    #   cohere
     #   datasets
     #   docker
     #   evaluate
@@ -1193,6 +1202,7 @@ tokenizers==0.22.2
     #   -c requirements/common.txt
     #   -r requirements/test/../common.txt
     #   -r requirements/test/rocm.in
+    #   cohere
     #   transformers
 torch-c-dlpack-ext==0.1.5
     # via tilelang
@@ -1243,6 +1253,8 @@ typer==0.24.1
     #   fastsafetensors
     #   perceptron
     #   transformers
+types-requests==2.33.0.20260712
+    # via cohere
 typing-extensions==4.15.0
     # via
     #   -c requirements/common.txt
@@ -1257,6 +1269,7 @@ typing-extensions==4.15.0
     #   azure-identity
     #   azure-storage-blob
     #   chz
+    #   cohere
     #   fastapi
     #   grpcio
     #   huggingface-hub
@@ -1302,6 +1315,7 @@ urllib3==2.6.3
     #   responses
     #   sentry-sdk
     #   tritonclient
+    #   types-requests
 uvicorn==0.42.0
     # via
     #   fastapi

tests/entrypoints/cohere/test_api_router.py

@@ -0,0 +1,425 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ``vllm/entrypoints/cohere/api_router.py``.
+
+Covers:
+
+* The optional-import guard: ``attach_router`` is a no-op when the
+  ``cohere`` SDK isn't installed.
+* The env-var opt-in gate: ``attach_router`` is a no-op unless
+  ``VLLM_ENABLE_COHERE_API=1`` is set.
+* The router wiring: response shapes (JSON + SSE), error translation,
+  and the ``cohere_serving_chat_v2 is None`` fallback (501 Not
+  Implemented).
+"""
+
+import json
+from argparse import Namespace
+from collections.abc import AsyncGenerator
+from http import HTTPStatus
+
+import pytest
+from fastapi import FastAPI, HTTPException
+from fastapi.exceptions import RequestValidationError
+from fastapi.testclient import TestClient
+
+from vllm.entrypoints.cohere import api_router as api_router_mod
+from vllm.entrypoints.cohere.api_router import attach_router
+from vllm.entrypoints.cohere.protocol import (
+    AssistantMessageResponse,
+    CohereChatV2Response,
+)
+from vllm.entrypoints.openai.engine.protocol import ErrorInfo, ErrorResponse
+from vllm.entrypoints.serve.utils.server_utils import (
+    http_exception_handler,
+    validation_exception_handler,
+)
+
+
+@pytest.fixture(autouse=True)
+def _enable_cohere_api(monkeypatch):
+    """Auto-enable the Cohere API gate for every test in this module.
+
+    The endpoint is opt-in in production (``VLLM_ENABLE_COHERE_API=1``);
+    every test in this file exercises the enabled path *except* the
+    dedicated gate test in :class:`TestEnvVarGate`, which unsets the
+    flag inside the test body.
+    """
+    monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "1")
+
+
+# ----------------------------------------------------------------------
+# Fakes
+# ----------------------------------------------------------------------
+
+
+class _Handler:
+    """Minimal stand-in for :class:`CohereServingChatV2` used by the
+    router. Each test sets ``self.result`` to either:
+
+    * a :class:`CohereChatV2Response` (non-streaming JSON path);
+    * an async generator yielding SSE frames (streaming path);
+    * an :class:`ErrorResponse` (error envelope path); or
+    * an exception (router-level 500 path).
+    """
+
+    def __init__(self, result):
+        self.result = result
+
+    async def create_chat_v2(self, request, raw_request):
+        if isinstance(self.result, Exception):
+            raise self.result
+        return self.result
+
+
+def _build_app(handler: _Handler | None) -> FastAPI:
+    app = FastAPI()
+    attach_router(app)
+    app.state.cohere_serving_chat_v2 = handler
+    return app
+
+
+def _build_app_with_vllm_handlers(handler: _Handler | None) -> FastAPI:
+    """Build a FastAPI app that mirrors the real vLLM setup by installing
+    ``validation_exception_handler`` and ``http_exception_handler``. The
+    :class:`CohereErrorEnvelopeMiddleware` registered by ``attach_router``
+    is expected to translate any resulting vLLM ``ErrorResponse`` body
+    into the ``CohereError`` wire shape.
+    """
+    app = FastAPI()
+    attach_router(app)
+    app.state.cohere_serving_chat_v2 = handler
+    # ``validation_exception_handler`` reads ``req.app.state.args``; the
+    # real cli builds this via argparse.
+    app.state.args = Namespace(log_error_stack=False)
+    app.exception_handler(RequestValidationError)(validation_exception_handler)
+    app.exception_handler(HTTPException)(http_exception_handler)
+    return app
+
+
+def _minimal_request_body() -> dict:
+    return {
+        "model": "m",
+        "messages": [{"role": "user", "content": "hi"}],
+    }
+
+
+# ----------------------------------------------------------------------
+# Optional-import guard
+# ----------------------------------------------------------------------
+
+
+class TestOptionalCohereImport:
+    """``attach_router`` probes for the SDK once at module load (because
+    the route handler uses types imported from ``cohere``) and stashes
+     the result in ``_SDK_AVAILABLE``. Tests simulate the "SDK missing"
+    state by flipping that flag for the duration of the test.
+
+    ``attach_router`` checks the env-var gate before the SDK probe, so
+    the SDK-missing branch is only reachable when the operator opts in
+    via ``VLLM_ENABLE_COHERE_API=1``. The flag-off-and-SDK-missing case
+    below exists to pin down that ordering — the flag-off short-circuits
+    """
+
+    def test_flag_off_and_sdk_missing_stays_silent_about_sdk(self, monkeypatch, caplog):
+        """Flag off doesn't do SDK-missing check.
+
+        When the operator hasn't opted in, ``attach_router`` must not
+        warn about the ``cohere`` SDK being missing: they never asked
+        for the endpoint, so surfacing the SDK gap is misleading noise.
+        Only the flag-off DEBUG message should fire.
+        """
+        monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False)
+        monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False)
+
+        with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
+            app = FastAPI()
+            attach_router(app)
+
+        paths = [getattr(r, "path", None) for r in app.routes]
+        assert "/cohere/v2/chat" not in paths
+        # The flag-off short-circuit ran; the SDK check never did.
+        assert not any(
+            "SDK is not installed" in rec.message for rec in caplog.records
+        ), "SDK-missing log leaked despite the flag being off"
+
+    def test_flag_on_but_sdk_missing_logs_warning(self, monkeypatch, caplog):
+        """Misconfiguration path: the operator explicitly opted into the
+        endpoint via ``VLLM_ENABLE_COHERE_API=1`` (already set by the
+        autouse fixture) but forgot to install ``cohere``.
+        """
+        monkeypatch.setattr(api_router_mod, "_SDK_AVAILABLE", False)
+
+        with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
+            app = FastAPI()
+            attach_router(app)
+
+        paths = [getattr(r, "path", None) for r in app.routes]
+        assert "/cohere/v2/chat" not in paths
+        warn_sdk_records = [
+            rec
+            for rec in caplog.records
+            if "VLLM_ENABLE_COHERE_API=1" in rec.message
+            and "SDK is not installed" in rec.message
+        ]
+        assert warn_sdk_records, (
+            "expected a WARNING that pairs the opt-in flag with the "
+            "missing SDK so operators notice the misconfiguration"
+        )
+        assert all(rec.levelname == "WARNING" for rec in warn_sdk_records)
+
+    def test_attach_router_registers_route_when_cohere_present(self):
+        app = _build_app(handler=None)
+        paths = [getattr(r, "path", None) for r in app.routes]
+        assert "/cohere/v2/chat" in paths
+
+
+# ----------------------------------------------------------------------
+# VLLM_ENABLE_COHERE_API gate
+# ----------------------------------------------------------------------
+
+
+class TestEnvVarGate:
+    """The Cohere v2 endpoint is opt-in via ``VLLM_ENABLE_COHERE_API``.
+
+    Even with the SDK installed, :func:`attach_router` must skip route
+    registration and middleware installation unless the env flag is
+    set. The autouse fixture on this module enables the flag by
+    default, so each test here explicitly disables it.
+    """
+
+    def test_attach_router_noop_when_flag_unset(self, monkeypatch, caplog):
+        monkeypatch.delenv("VLLM_ENABLE_COHERE_API", raising=False)
+
+        # The flag-off skip logs at DEBUG on purpose: this is the default
+        # state for every non-Cohere vLLM deployment, so an INFO log on
+        # every server startup would be pointless noise. The test raises
+        # caplog's level accordingly.
+        with caplog.at_level("DEBUG", logger="vllm.entrypoints.cohere.api_router"):
+            app = FastAPI()
+            attach_router(app)
+
+        paths = [getattr(r, "path", None) for r in app.routes]
+        assert "/cohere/v2/chat" not in paths
+        debug_flag_records = [
+            rec
+            for rec in caplog.records
+            if "VLLM_ENABLE_COHERE_API is not set" in rec.message
+        ]
+        assert debug_flag_records, (
+            "expected a DEBUG message that the cohere flag is off"
+        )
+        assert all(rec.levelname == "DEBUG" for rec in debug_flag_records)
+
+    def test_attach_router_noop_when_flag_zero(self, monkeypatch):
+        monkeypatch.setenv("VLLM_ENABLE_COHERE_API", "0")
+
+        app = FastAPI()
+        attach_router(app)
+
+        paths = [getattr(r, "path", None) for r in app.routes]
+        assert "/cohere/v2/chat" not in paths
+
+
+# ----------------------------------------------------------------------
+# Endpoint behavior
+# ----------------------------------------------------------------------
+
+
+class TestEndpoint:
+    def test_501_when_handler_missing(self):
+        app = _build_app(handler=None)
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json=_minimal_request_body())
+        assert r.status_code == HTTPStatus.NOT_IMPLEMENTED
+        body = r.json()
+        assert "does not support" in body["message"]
+        assert "id" not in body  # excluded by ``exclude_none=True``
+
+    def test_non_streaming_response_is_json(self):
+        msg = AssistantMessageResponse(content=[{"type": "text", "text": "hello"}])
+        result = CohereChatV2Response(id="r1", finish_reason="COMPLETE", message=msg)
+        app = _build_app(handler=_Handler(result))
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json=_minimal_request_body())
+        assert r.status_code == HTTPStatus.OK
+        assert r.headers["content-type"].startswith("application/json")
+        body = r.json()
+        assert body["id"] == "r1"
+        assert body["finish_reason"] == "COMPLETE"
+        assert body["message"]["content"][0]["text"] == "hello"
+
+    def test_streaming_response_is_sse(self):
+        async def _gen() -> AsyncGenerator[str, None]:
+            yield 'data: {"type":"message-start"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        app = _build_app(handler=_Handler(_gen()))
+        with TestClient(app) as client:
+            r = client.post(
+                "/cohere/v2/chat",
+                json={**_minimal_request_body(), "stream": True},
+            )
+        assert r.status_code == HTTPStatus.OK
+        assert r.headers["content-type"].startswith("text/event-stream")
+        body = r.text
+        assert "message-start" in body
+        assert body.rstrip().endswith("[DONE]")
+
+    def test_error_response_translated_to_cohere_envelope(self):
+        err = ErrorResponse(
+            error=ErrorInfo(
+                message="bad request",
+                type="bad_request",
+                code=400,
+            )
+        )
+        app = _build_app(handler=_Handler(err))
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json=_minimal_request_body())
+        assert r.status_code == HTTPStatus.BAD_REQUEST
+        body = r.json()
+        assert body == {"message": "bad request"}
+
+    def test_handler_exception_returns_500_envelope(self):
+        app = _build_app(handler=_Handler(RuntimeError("kaboom")))
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json=_minimal_request_body())
+        assert r.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
+        body = r.json()
+        assert body == {"message": "kaboom"}
+
+    def test_non_json_content_type_rejected(self):
+        """The ``validate_json_request`` dependency raises
+        ``RequestValidationError`` (HTTP 422) for non-JSON content
+        types, matching the behavior of the other vLLM API routers.
+        """
+        app = _build_app(handler=None)
+        with TestClient(app) as client:
+            r = client.post(
+                "/cohere/v2/chat",
+                content=json.dumps(_minimal_request_body()),
+                headers={"content-type": "text/plain"},
+            )
+        assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
+
+    def test_invalid_body_returns_422(self):
+        # ``model`` is required; omit it to trip Pydantic validation.
+        app = _build_app(handler=None)
+        with TestClient(app) as client:
+            r = client.post(
+                "/cohere/v2/chat",
+                json={"messages": [{"role": "user", "content": "hi"}]},
+            )
+        assert r.status_code == HTTPStatus.UNPROCESSABLE_ENTITY
+
+
+# ----------------------------------------------------------------------
+# CohereErrorEnvelopeMiddleware
+# ----------------------------------------------------------------------
+
+
+class TestCohereErrorEnvelope:
+    """When the app installs vLLM's global exception handlers, validation
+    and HTTP errors escape as ``ErrorResponse`` bodies before the route
+    handler runs. The middleware installed by ``attach_router`` must
+    normalise those bodies to the ``CohereError`` shape declared on the
+    endpoint's OpenAPI ``responses`` map so schema-conformance tests
+    (``test_openai_schema.py``) don't see a mismatch on ``/cohere/*``
+    responses.
+    """
+
+    def test_validation_error_body_is_cohere_shaped(self):
+        # ``model=""`` and ``messages=[]`` trip our custom field
+        # validators, which raise pydantic ValueErrors and are routed
+        # through ``validation_exception_handler`` in the real vLLM
+        # server (producing the ``{"error": {...}}`` shape).
+        app = _build_app_with_vllm_handlers(handler=None)
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json={"messages": [], "model": ""})
+        assert r.status_code == HTTPStatus.BAD_REQUEST
+        body = r.json()
+        # ``CohereError`` has ``message`` at the top level, not nested
+        # under an ``error`` envelope.
+        assert "error" not in body
+        assert "message" in body
+        assert isinstance(body["message"], str) and body["message"]
+
+    def test_http_error_body_is_cohere_shaped(self):
+        # A raised ``HTTPException`` from anywhere in the request cycle
+        # is routed through ``http_exception_handler`` (producing the
+        # ``ErrorResponse`` shape) and must be translated.
+        app = _build_app_with_vllm_handlers(handler=None)
+
+        @app.get("/cohere/v2/boom")
+        async def _boom():
+            raise HTTPException(status_code=418, detail="teapot")
+
+        with TestClient(app) as client:
+            r = client.get("/cohere/v2/boom")
+        assert r.status_code == 418
+        body = r.json()
+        assert body == {"message": "teapot"}
+
+    def test_non_cohere_path_is_not_translated(self):
+        app = _build_app_with_vllm_handlers(handler=None)
+
+        @app.get("/v1/other")
+        async def _other():
+            raise HTTPException(status_code=400, detail="nope")
+
+        with TestClient(app) as client:
+            r = client.get("/v1/other")
+        assert r.status_code == HTTPStatus.BAD_REQUEST
+        body = r.json()
+        # Non-cohere paths keep the vLLM ``ErrorResponse`` shape.
+        assert "error" in body
+        assert body["error"]["message"] == "nope"
+
+    def test_streaming_response_passes_through(self):
+        # SSE responses have content-type text/event-stream; the
+        # middleware must never buffer these (which would break
+        # streaming) even though they're on ``/cohere/*``.
+        async def _gen() -> AsyncGenerator[str, None]:
+            yield 'data: {"type":"message-start"}\n\n'
+            yield "data: [DONE]\n\n"
+
+        app = _build_app_with_vllm_handlers(handler=_Handler(_gen()))
+        with TestClient(app) as client:
+            r = client.post(
+                "/cohere/v2/chat",
+                json={**_minimal_request_body(), "stream": True},
+            )
+        assert r.status_code == HTTPStatus.OK
+        assert r.headers["content-type"].startswith("text/event-stream")
+        assert "message-start" in r.text
+        assert r.text.rstrip().endswith("[DONE]")
+
+    def test_already_cohere_shaped_body_passes_through(self):
+        # When the handler returns an ``ErrorResponse`` the route
+        # itself translates it to ``CohereError``; the middleware sees
+        # the ``CohereError`` shape and must leave it alone.
+        err = ErrorResponse(
+            error=ErrorInfo(message="already cohere", type="Bad Request", code=400)
+        )
+        app = _build_app_with_vllm_handlers(handler=_Handler(err))
+        with TestClient(app) as client:
+            r = client.post("/cohere/v2/chat", json=_minimal_request_body())
+        assert r.status_code == HTTPStatus.BAD_REQUEST
+        body = r.json()
+        # No ``error`` wrapper: the route already emitted the wire shape.
+        assert body == {"message": "already cohere"}
+
+    def test_request_id_preserved_in_translated_body(self):
+        # Client-provided ``X-Request-Id`` should be echoed as
+        # ``CohereError.id`` so callers can correlate failures.
+        app = _build_app_with_vllm_handlers(handler=None)
+        with TestClient(app) as client:
+            r = client.post(
+                "/cohere/v2/chat",
+                json={"messages": [], "model": ""},
+                headers={"X-Request-Id": "req-abc"},
+            )
+        assert r.status_code == HTTPStatus.BAD_REQUEST
+        body = r.json()
+        assert body.get("id") == "req-abc"

tests/entrypoints/cohere/test_chat_v2.py

@@ -0,0 +1,259 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""End-to-end integration tests for the ``POST /cohere/v2/chat`` endpoint.
+
+These tests spin up a real ``vllm serve`` process via
+:class:`tests.utils.RemoteOpenAIServer` and exercise the Cohere Chat v2
+contract over HTTP. They mirror the pattern used by
+:mod:`tests.entrypoints.anthropic.test_messages`.
+
+Two layers of integration are covered:
+
+1. Raw HTTP via :mod:`httpx` — always runs, verifies the wire contract.
+2. Cohere SDK (``pip install cohere``) — auto-skipped when the optional
+   dependency isn't installed, verifies SDK-level interop.
+
+We use a tiny generic chat model (``HuggingFaceTB/SmolLM2-135M-Instruct``,
+≈135M params) so the fixture stays runnable on CPU-only laptops. The
+Cohere v2 endpoint is model-agnostic and just performs the v2 ↔ OpenAI
+chat-completion translation, so the choice of model only matters for
+response *shape* — the translation logic itself is unit-tested
+elsewhere.
+"""
+
+import json
+
+import httpx
+import pytest
+import pytest_asyncio
+
+from tests.utils import RemoteOpenAIServer
+
+# Tiny chat-tuned model so the fixture is cheap to boot on CPU-only
+# hosts (Mac, CI runners without GPUs). The Cohere v2 translation is
+# completely independent of which underlying chat model is loaded.
+MODEL_NAME = "HuggingFaceTB/SmolLM2-135M-Instruct"
+SERVED_MODEL_NAME = "command-r-plus-08-2024"
+
+
+@pytest.fixture(scope="module")
+def server():
+    args = [
+        # 1k tokens is plenty for the prompts these tests send; trimming
+        # it from the default keeps the KV-cache footprint tiny.
+        "--max-model-len",
+        "1024",
+        "--max-num-seqs",
+        "4",
+        "--dtype",
+        "bfloat16",
+        "--enforce-eager",
+        # Advertise a Cohere model name so Cohere SDK ``model=`` calls
+        # round-trip without the ``model not found`` check.
+        "--served-model-name",
+        SERVED_MODEL_NAME,
+        # Disable the reasoning-model path so the test doesn't require
+        # the conversation to surface a thinking block (SmolLM2 doesn't
+        # emit Cohere-style reasoning tokens out of the box).
+        "--no-cohere-is-reasoning-model",
+    ]
+
+    # Cap the CPU KV-cache pool the vLLM CPU backend reserves at
+    # startup. The default (4 GB) trips Mac's RAM check on smaller
+    # machines; 1 GB is more than enough for max_model_len=1024 and
+    # max_num_seqs=4.
+    # ``VLLM_ENABLE_COHERE_API=1`` flips the opt-in gate that
+    # ``vllm.entrypoints.cohere.api_router.attach_router`` reads to
+    # decide whether to register ``POST /cohere/v2/chat`` at server
+    # startup. Without it the fixture would boot a server that returns
+    # 404 for every request the tests below make.
+    env_dict = {
+        "VLLM_CPU_KVCACHE_SPACE": "1",
+        "VLLM_ENABLE_COHERE_API": "1",
+    }
+
+    with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
+        yield remote_server
+
+
+# ----------------------------------------------------------------------
+# Layer 1: raw HTTP contract (no optional deps)
+# ----------------------------------------------------------------------
+
+
+@pytest_asyncio.fixture
+async def httpx_client(server):
+    async with httpx.AsyncClient(
+        base_url=server.url_root, timeout=httpx.Timeout(120.0)
+    ) as client:
+        yield client
+
+
+@pytest.mark.asyncio
+async def test_cohere_v2_chat_non_streaming(httpx_client: httpx.AsyncClient):
+    resp = await httpx_client.post(
+        "/cohere/v2/chat",
+        json={
+            "model": SERVED_MODEL_NAME,
+            "messages": [{"role": "user", "content": "Say hi."}],
+            "max_tokens": 16,
+            "stream": False,
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    payload = resp.json()
+    # The response envelope follows Cohere v2's schema.
+    assert "message" in payload
+    assert payload["message"]["role"] == "assistant"
+    # ``content`` is a list of content blocks; the synthesized
+    # ``CohereServingChatV2`` should always emit at least one ``text``
+    # block (it falls back to an empty block when the model returned
+    # nothing).
+    content = payload["message"]["content"]
+    assert isinstance(content, list) and len(content) >= 1
+    assert content[0]["type"] == "text"
+    # ``usage`` is always populated by the translator.
+    assert "usage" in payload
+    assert "billed_units" in payload["usage"]
+    assert "tokens" in payload["usage"]
+    # ``finish_reason`` is one of Cohere's enum values.
+    assert payload["finish_reason"] in {
+        "COMPLETE",
+        "MAX_TOKENS",
+        "STOP_SEQUENCE",
+        "TOOL_CALL",
+        "ERROR",
+    }
+
+
+@pytest.mark.asyncio
+async def test_cohere_v2_chat_streaming(httpx_client: httpx.AsyncClient):
+    """Streaming returns SSE frames in the v2 message-lifecycle shape."""
+    events: list[dict] = []
+    async with httpx_client.stream(
+        "POST",
+        "/cohere/v2/chat",
+        json={
+            "model": SERVED_MODEL_NAME,
+            "messages": [{"role": "user", "content": "Hi"}],
+            "max_tokens": 8,
+            "stream": True,
+        },
+    ) as resp:
+        assert resp.status_code == 200, await resp.aread()
+        assert resp.headers["content-type"].startswith("text/event-stream")
+        async for line in resp.aiter_lines():
+            if not line.startswith("data: "):
+                continue
+            data = line[len("data: ") :]
+            if data == "[DONE]":
+                events.append({"type": "_DONE_"})
+                continue
+            events.append(json.loads(data))
+
+    types = [ev["type"] for ev in events]
+    # The lifecycle always starts with message-start and ends with [DONE]
+    # preceded by message-end.
+    assert types[0] == "message-start"
+    assert types[-1] == "_DONE_"
+    assert types[-2] == "message-end"
+    # message-start must carry the chunk/message id.
+    assert events[0].get("id")
+
+
+@pytest.mark.asyncio
+async def test_cohere_v2_chat_validation_error_returns_400(
+    httpx_client: httpx.AsyncClient,
+):
+    # Missing required ``model`` field → FastAPI/Pydantic returns 400.
+    resp = await httpx_client.post(
+        "/cohere/v2/chat",
+        json={"messages": []},
+    )
+    assert resp.status_code == 400
+
+
+@pytest.mark.asyncio
+async def test_cohere_v2_chat_documents_field_accepted(
+    httpx_client: httpx.AsyncClient,
+):
+    """The v2 endpoint forwards ``documents`` into chat_template_kwargs.
+
+    We only assert the request is accepted and produces a 200 response —
+    the renderer-level effect is covered by ``tests/renderers/test_cohere.py``.
+    """
+    resp = await httpx_client.post(
+        "/cohere/v2/chat",
+        json={
+            "model": SERVED_MODEL_NAME,
+            "messages": [{"role": "user", "content": "Summarize."}],
+            "documents": [{"id": "d1", "data": {"title": "T", "snippet": "S"}}],
+            "max_tokens": 16,
+            "stream": False,
+        },
+    )
+    assert resp.status_code == 200, resp.text
+
+
+# ----------------------------------------------------------------------
+# Layer 2: Cohere SDK round-trip (auto-skipped if SDK absent)
+# ----------------------------------------------------------------------
+
+
+@pytest_asyncio.fixture
+async def cohere_async_client(server):
+    cohere = pytest.importorskip("cohere")
+    # The vLLM endpoint is mounted at ``/cohere/v2/chat`` while the
+    # cohere SDK targets ``${base_url}/v2/chat``; point base_url at the
+    # ``/cohere`` prefix so paths line up.
+    client = cohere.AsyncClientV2(
+        api_key="dummy",
+        base_url=server.url_for("cohere"),
+    )
+    try:
+        yield client
+    finally:
+        # ``AsyncClientV2`` exposes a sync close; if a future version
+        # adds aclose we still close cleanly.
+        close = getattr(client, "aclose", None) or getattr(client, "close", None)
+        if close is not None:
+            result = close()
+            if hasattr(result, "__await__"):
+                await result
+
+
+@pytest.mark.asyncio
+async def test_cohere_sdk_non_streaming(cohere_async_client):
+    resp = await cohere_async_client.chat(
+        model=SERVED_MODEL_NAME,
+        messages=[{"role": "user", "content": "Say hi."}],
+        max_tokens=16,
+    )
+    # SDK parses our JSON into typed objects.
+    assert resp.message.role == "assistant"
+    assert resp.message.content is not None
+    assert len(resp.message.content) >= 1
+    assert resp.message.content[0].type == "text"
+    assert resp.finish_reason in {
+        "COMPLETE",
+        "MAX_TOKENS",
+        "STOP_SEQUENCE",
+        "TOOL_CALL",
+        "ERROR",
+    }
+
+
+@pytest.mark.asyncio
+async def test_cohere_sdk_streaming(cohere_async_client):
+    events: list[str] = []
+    stream = cohere_async_client.chat_stream(
+        model=SERVED_MODEL_NAME,
+        messages=[{"role": "user", "content": "Hi"}],
+        max_tokens=8,
+    )
+    async for ev in stream:
+        events.append(ev.type)
+
+    assert events, "SDK stream yielded no events"
+    assert events[0] == "message-start"
+    assert events[-1] == "message-end"

tests/entrypoints/cohere/test_cohere_chat_message.py

@@ -0,0 +1,135 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for :mod:`vllm.entrypoints.cohere.cohere_chat_message`.
+
+The module hosts the citation-carrying subclasses of the OpenAI chat
+completion protocol (``CohereChatMessage`` / ``CohereDeltaMessage``) plus
+the shared ``Citation`` / ``CitationSource`` shapes. These tests pin:
+
+1. Serialization behavior we own (empty-``citations`` cleanup on the
+   subclass ``_serialize``, ``exclude_none`` shape, round-trip).
+2. The constraint that ``type`` discriminators reject unknown values.
+"""
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.entrypoints.cohere.cohere_chat_message import (
+    Citation,
+    CitationSource,
+    CohereChatMessage,
+    CohereDeltaMessage,
+)
+
+# ======================================================================
+# CitationSource
+# ======================================================================
+
+
+class TestCitationSource:
+    def test_invalid_type_rejected(self):
+        with pytest.raises(ValidationError):
+            CitationSource(type="other")  # type: ignore[arg-type]
+
+    def test_none_fields_excluded_from_dump(self):
+        s = CitationSource(type="document", id="d1")
+        assert s.model_dump(exclude_none=True) == {
+            "type": "document",
+            "id": "d1",
+        }
+
+    def test_resolved_source_survives_json_round_trip(self):
+        # The parser produces fully-resolved sources (see
+        # ``_melody_sources_to_vllm`` in
+        # ``vllm/reasoning/cohere_command_reasoning_parser.py``) which
+        # then flow through the internal streaming pipeline (parser ->
+        # OpenAI stream chunk -> ``_chat_completion_stream_to_v2``).
+        # Every field must round-trip through ``model_dump_json`` at
+        # that layer or the wire event ends up missing pieces.
+        src = CitationSource(
+            type="tool",
+            id="res_a0",
+            tool_output={"id": "res_a0", "text": "r"},
+        )
+        src2 = CitationSource.model_validate_json(src.model_dump_json())
+        assert src2.type == "tool"
+        assert src2.id == "res_a0"
+        assert src2.tool_output == {"id": "res_a0", "text": "r"}
+
+
+# ======================================================================
+# Citation
+# ======================================================================
+
+
+class TestCitation:
+    def test_invalid_type_rejected(self):
+        with pytest.raises(ValidationError):
+            Citation(type="OTHER")  # type: ignore[arg-type]
+
+    def test_dump_excludes_none_fields(self):
+        c = Citation(start=0, end=5, text="hello")
+        dumped = c.model_dump(exclude_none=True)
+        assert dumped == {
+            "start": 0,
+            "end": 5,
+            "text": "hello",
+            "sources": [],
+        }
+
+
+# ======================================================================
+# CohereDeltaMessage
+# ======================================================================
+
+
+class TestCohereDeltaMessage:
+    def test_citations_omitted_from_dump_when_none(self):
+        # Non-grounded deltas should never carry ``citations: null`` on
+        # the wire even when routed through the Cohere subclass.
+        d = CohereDeltaMessage(content="hello")
+        dumped = d.model_dump(exclude_none=True)
+        assert "citations" not in dumped
+
+    def test_citations_dump_round_trip(self):
+        d = CohereDeltaMessage(citations=[Citation(start=0, end=5, text="hi")])
+        dumped = d.model_dump(exclude_none=True)
+        assert dumped["citations"][0]["text"] == "hi"
+        # Round-trip through the subclass to confirm the field validates.
+        d2 = CohereDeltaMessage.model_validate(dumped)
+        assert d2.citations[0].text == "hi"
+
+
+# ======================================================================
+# CohereChatMessage
+# ======================================================================
+
+
+class TestCohereChatMessage:
+    def test_citations_omitted_from_dump_when_none(self):
+        m = CohereChatMessage(role="assistant", content="hi")
+        dumped = m.model_dump(exclude_none=True)
+        assert "citations" not in dumped
+
+    def test_citations_dump_round_trip(self):
+        m = CohereChatMessage(
+            role="assistant",
+            content="hello",
+            citations=[
+                Citation(
+                    start=0,
+                    end=5,
+                    text="hello",
+                    sources=[CitationSource(type="document", id="d1")],
+                )
+            ],
+        )
+        dumped = m.model_dump(exclude_none=True)
+        assert dumped["role"] == "assistant"
+        assert dumped["content"] == "hello"
+        assert dumped["citations"][0]["text"] == "hello"
+        assert dumped["citations"][0]["sources"][0]["id"] == "d1"
+        # Round-trip through the subclass to confirm the field validates.
+        m2 = CohereChatMessage.model_validate(dumped)
+        assert m2.citations[0].text == "hello"
+        assert m2.citations[0].sources[0].id == "d1"

tests/entrypoints/cohere/test_protocol.py

@@ -0,0 +1,217 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for ``vllm/entrypoints/cohere/protocol.py``.
+
+The module is mostly a thin wrapper around the official ``cohere`` SDK
+types plus a few local additions:
+
+* :class:`CohereError` envelope.
+* :class:`CohereChatV2Request` (model required, ``max_tokens`` non-negative).
+* :class:`CohereChatV2Response` plus the usage / logprob helpers.
+* The streaming event subclasses that bake a wire-format ``type``
+  discriminator into ``model_dump()`` so SSE consumers can demux on it.
+"""
+
+import pytest
+from pydantic import ValidationError
+
+from vllm.entrypoints.cohere.protocol import (
+    AssistantMessageResponse,
+    CitationEndEvent,
+    CitationStartEvent,
+    CohereChatV2Request,
+    CohereChatV2Response,
+    CohereError,
+    CohereUsageTokens,
+    ContentDeltaEvent,
+    ContentEndEvent,
+    ContentStartEvent,
+    MessageEndEvent,
+    MessageStartEvent,
+    ToolCallDeltaEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+    ToolPlanDeltaEvent,
+)
+
+# ======================================================================
+# CohereError
+# ======================================================================
+
+
+class TestCohereError:
+    def test_model_dump_excludes_none(self):
+        err = CohereError(message="boom")
+        assert err.model_dump(exclude_none=True) == {"message": "boom"}
+
+
+# ======================================================================
+# CohereChatV2Request
+# ======================================================================
+
+
+class TestCohereChatV2Request:
+    def test_empty_model_rejected(self):
+        with pytest.raises(ValidationError, match="model is required"):
+            CohereChatV2Request(model="", messages=[{"role": "user", "content": "hi"}])
+
+    def test_negative_max_tokens_rejected(self):
+        with pytest.raises(ValidationError, match="non-negative"):
+            CohereChatV2Request(
+                model="m",
+                messages=[{"role": "user", "content": "hi"}],
+                max_tokens=-1,
+            )
+
+    def test_zero_max_tokens_allowed(self):
+        # Zero is allowed (the docs allow 0 -> return prompt only).
+        req = CohereChatV2Request(
+            model="m",
+            messages=[{"role": "user", "content": "hi"}],
+            max_tokens=0,
+        )
+        assert req.max_tokens == 0
+
+    def test_invalid_tool_choice_rejected(self):
+        with pytest.raises(ValidationError):
+            CohereChatV2Request(
+                model="m",
+                messages=[{"role": "user", "content": "hi"}],
+                tool_choice="ANY",  # not REQUIRED/NONE
+            )
+
+
+# ======================================================================
+# Usage / Logprobs
+# ======================================================================
+
+
+class TestUsage:
+    def test_tokens_serialization(self):
+        u = CohereUsageTokens(input_tokens=10, output_tokens=5)
+        assert u.model_dump() == {"input_tokens": 10.0, "output_tokens": 5.0}
+
+
+# ======================================================================
+# CohereChatV2Response
+# ======================================================================
+
+
+class TestCohereChatV2Response:
+    def test_invalid_finish_reason_rejected(self):
+        msg = AssistantMessageResponse(content=[{"type": "text", "text": "hi"}])
+        with pytest.raises(ValidationError):
+            CohereChatV2Response(
+                id="r1",
+                finish_reason="NOT_A_REASON",  # type: ignore[arg-type]
+                message=msg,
+            )
+
+
+# ======================================================================
+# Streaming event ``type`` discriminator baked into model_dump()
+# ======================================================================
+
+
+class TestStreamingEventTypeField:
+    """Each event subclass adds a ``type: Literal[...]`` field with a
+    default so ``model_dump()`` always emits the wire-format discriminator
+    (the parent SDK classes don't declare ``type`` as a Pydantic field).
+    """
+
+    @pytest.mark.parametrize(
+        "cls, expected_type, kwargs",
+        [
+            (
+                MessageStartEvent,
+                "message-start",
+                {"id": "a", "delta": {"message": {"role": "assistant"}}},
+            ),
+            (
+                ContentStartEvent,
+                "content-start",
+                {
+                    "index": 0,
+                    "delta": {"message": {"content": {"type": "text", "text": ""}}},
+                },
+            ),
+            (
+                ContentDeltaEvent,
+                "content-delta",
+                {
+                    "index": 0,
+                    "delta": {"message": {"content": {"text": "hi"}}},
+                },
+            ),
+            (ContentEndEvent, "content-end", {"index": 0}),
+            (
+                ToolPlanDeltaEvent,
+                "tool-plan-delta",
+                {"delta": {"message": {"tool_plan": "thinking"}}},
+            ),
+            (
+                ToolCallStartEvent,
+                "tool-call-start",
+                {
+                    "index": 0,
+                    "delta": {
+                        "message": {
+                            "tool_calls": {
+                                "id": "c1",
+                                "type": "function",
+                                "function": {"name": "f", "arguments": ""},
+                            }
+                        }
+                    },
+                },
+            ),
+            (
+                ToolCallDeltaEvent,
+                "tool-call-delta",
+                {
+                    "index": 0,
+                    "delta": {
+                        "message": {"tool_calls": {"function": {"arguments": "{}"}}}
+                    },
+                },
+            ),
+            (ToolCallEndEvent, "tool-call-end", {"index": 0}),
+            (
+                CitationStartEvent,
+                "citation-start",
+                {
+                    "index": 0,
+                    "delta": {
+                        "message": {
+                            "citations": {"start": 0, "end": 5, "text": "hello"}
+                        }
+                    },
+                },
+            ),
+            (CitationEndEvent, "citation-end", {"index": 0}),
+            (
+                MessageEndEvent,
+                "message-end",
+                {"id": "a", "delta": {"finish_reason": "COMPLETE"}},
+            ),
+        ],
+    )
+    def test_type_field_default(self, cls, expected_type, kwargs):
+        ev = cls(**kwargs)
+        # type field is auto-populated from the Literal default.
+        assert ev.type == expected_type
+        # The discriminator must be present in the serialized payload so
+        # clients reading the stream can demux on it.
+        dumped = ev.model_dump(exclude_none=True)
+        assert dumped["type"] == expected_type
+        # Same in JSON form (what ``_emit`` serializes).
+        assert f'"type":"{expected_type}"' in ev.model_dump_json(exclude_none=True)
+
+    def test_type_field_cannot_be_overridden_to_wrong_value(self):
+        # Literal types reject any value other than the bake-in default.
+        with pytest.raises(ValidationError):
+            MessageStartEvent(
+                id="a",
+                delta={"message": {"role": "assistant"}},
+                type="other",  # type: ignore[arg-type]
+            )

tests/entrypoints/cohere/test_registry_and_args.py

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Sanity tests for the small registry / CLI-arg / config additions made
+by the Cohere v2 chat API:
+
+* ``vllm/renderers/registry.py``: new ``"cohere"`` renderer entry.
+* ``vllm/tokenizers/registry.py``: new ``"cohere"`` tokenizer entry
+  (aliased to the cached HF tokenizer).
+* ``vllm/entrypoints/openai/cli_args.py``: new
+  ``cohere_is_reasoning_model`` field.
+* ``vllm/config/model.py``: ``"cohere"`` added to ``TokenizerMode``
+  Literal.
+"""
+
+import dataclasses
+import typing
+
+import pytest
+
+from vllm.config.model import TokenizerMode
+from vllm.entrypoints.openai.cli_args import BaseFrontendArgs, make_arg_parser
+from vllm.renderers.registry import RENDERER_REGISTRY
+from vllm.tokenizers.registry import TokenizerRegistry
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+
+class TestRendererRegistry:
+    def test_cohere_renderer_registered(self):
+        # The registry resolves to the importable ``CohereRenderer`` class.
+        cls = RENDERER_REGISTRY.load_renderer_cls("cohere")
+        assert cls.__name__ == "CohereRenderer"
+        # Sanity: class lives in the cohere renderer module.
+        assert cls.__module__ == "vllm.renderers.cohere"
+
+
+class TestTokenizerRegistry:
+    def test_cohere_aliased_to_cached_hf_tokenizer(self):
+        # ``cohere`` mode uses the standard HF tokenizer; only the
+        # renderer stage is replaced. This test guards against accidental
+        # divergence.
+        cls = TokenizerRegistry.load_tokenizer_cls("cohere")
+        assert cls.__name__ == "CachedHfTokenizer"
+
+
+class TestTokenizerModeLiteral:
+    def test_cohere_is_a_valid_tokenizer_mode(self):
+        # The Literal must enumerate ``"cohere"`` so engine arg parsing
+        # accepts ``--tokenizer-mode cohere``.
+        modes = typing.get_args(TokenizerMode)
+        assert "cohere" in modes
+
+
+# ----------------------------------------------------------------------
+# ``--cohere-is-reasoning-model`` CLI flag
+# ----------------------------------------------------------------------
+
+
+class TestCohereCliArg:
+    """Verifies the new ``--cohere-is-reasoning-model`` flag end-to-end
+    through ``make_arg_parser`` — mirrors the pattern used in
+    :mod:`tests.entrypoints.openai.test_cli_args`.
+    """
+
+    def test_default_value_on_dataclass_is_true(self):
+        fields = {f.name: f for f in dataclasses.fields(BaseFrontendArgs)}
+        assert "cohere_is_reasoning_model" in fields
+        field = fields["cohere_is_reasoning_model"]
+        assert field.default is True
+        assert field.type is bool
+
+    @pytest.fixture
+    def serve_parser(self) -> FlexibleArgumentParser:
+        parser = FlexibleArgumentParser()
+        return make_arg_parser(parser)
+
+    def test_default_via_argparse_is_true(self, serve_parser: FlexibleArgumentParser):
+        # No flag supplied → dataclass default (True) wins.
+        args = serve_parser.parse_args(["--model", "m"])
+        assert args.cohere_is_reasoning_model is True
+
+    def test_explicit_false_via_argparse(self, serve_parser: FlexibleArgumentParser):
+        # Boolean dataclass fields are wired up as ``--flag value`` /
+        # ``--no-flag`` pairs by FlexibleArgumentParser.
+        args = serve_parser.parse_args(
+            ["--model", "m", "--no-cohere-is-reasoning-model"]
+        )
+        assert args.cohere_is_reasoning_model is False
+
+    def test_explicit_true_via_argparse(self, serve_parser: FlexibleArgumentParser):
+        args = serve_parser.parse_args(["--model", "m", "--cohere-is-reasoning-model"])
+        assert args.cohere_is_reasoning_model is True

tests/entrypoints/cohere/test_serving_streaming.py

@@ -0,0 +1,944 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for the Cohere v2 SSE stream conversion in
+``vllm/entrypoints/cohere/serving.py``.
+
+The stream-translation entry point is
+:meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an
+async iterable of OpenAI SSE chunks into Cohere's
+``message-start → (content|tool-call|citation)* → message-end → [DONE]``
+event stream.
+
+We test the helpers (``_StreamState``, ``_handle_*_delta``, etc.) in
+isolation, plus a handful of end-to-end scenarios that exercise the
+state machine.
+"""
+
+import json
+from collections.abc import AsyncGenerator
+from typing import Any
+
+import pytest
+
+from vllm.entrypoints.cohere.cohere_chat_message import (
+    Citation as VLLMCitation,
+)
+from vllm.entrypoints.cohere.cohere_chat_message import (
+    CitationSource,
+)
+from vllm.entrypoints.cohere.protocol import (
+    CohereChatV2Request,
+    MessageStartEvent,
+)
+from vllm.entrypoints.cohere.serving import (
+    _DONE_FRAME,
+    CohereServingChatV2,
+    ContentBlockType,
+    _emit,
+    _sse,
+    _StreamState,
+)
+from vllm.entrypoints.openai.chat_completion.protocol import (
+    ChatCompletionStreamResponse,
+)
+
+# ----------------------------------------------------------------------
+# Helpers
+# ----------------------------------------------------------------------
+
+
+class _FakeServing(CohereServingChatV2):
+    """Lightweight stand-in for :class:`CohereServingChatV2` that skips
+    the heavy ``OpenAIServingChat.__init__`` chain (which would need a
+    real engine client, model registry, render service, etc.).
+
+    Only ``_is_reasoning_model`` is read by the methods under test
+    (``_chat_completion_stream_to_v2`` and the per-delta handlers); the
+    rest is dead weight for unit testing.
+    """
+
+    def __init__(self, is_reasoning_model: bool = True) -> None:
+        # Intentionally skipping super().__init__ — see class docstring.
+        self._is_reasoning_model = is_reasoning_model
+
+
+def _serving(is_reasoning_model: bool = True) -> CohereServingChatV2:
+    return _FakeServing(is_reasoning_model=is_reasoning_model)
+
+
+def _parse_event(frame: str) -> dict[str, Any]:
+    """Strip the ``data: ... \\n\\n`` wrapper and parse the JSON payload."""
+    assert frame.startswith("data: ")
+    assert frame.endswith("\n\n")
+    return json.loads(frame[len("data: ") : -2])
+
+
+def _make_chunk(
+    *,
+    chunk_id: str = "chunk_0",
+    role: str | None = None,
+    content: str | None = None,
+    reasoning: str | None = None,
+    tool_calls: list[dict[str, Any]] | None = None,
+    finish_reason: str | None = None,
+    usage: dict[str, Any] | None = None,
+    omit_choices: bool = False,
+    citations: list[Any] | None = None,
+) -> str:
+    """Build the ``data: {...}\\n\\n`` SSE frame the production code
+    consumes."""
+    delta: dict[str, Any] = {}
+    if role is not None:
+        delta["role"] = role
+    if content is not None:
+        delta["content"] = content
+    if reasoning is not None:
+        delta["reasoning"] = reasoning
+    if tool_calls is not None:
+        delta["tool_calls"] = tool_calls
+    if citations is not None:
+        delta["citations"] = citations
+
+    payload: dict[str, Any] = {
+        "id": chunk_id,
+        "object": "chat.completion.chunk",
+        "created": 0,
+        "model": "m",
+    }
+    if not omit_choices:
+        payload["choices"] = [
+            {"index": 0, "delta": delta, "finish_reason": finish_reason}
+        ]
+    else:
+        payload["choices"] = []
+    if usage is not None:
+        payload["usage"] = usage
+    # Use ChatCompletionStreamResponse to normalize the payload shape.
+    chunk = ChatCompletionStreamResponse.model_validate(payload)
+    return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n"
+
+
+def _make_done() -> str:
+    return "data: [DONE]\n\n"
+
+
+async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]:
+    for item in items:
+        yield item
+
+
+async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]:
+    """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect
+    the emitted SSE frames."""
+    request = CohereChatV2Request(
+        model="m", messages=[{"role": "user", "content": "hi"}], stream=True
+    )
+    gen = serving._chat_completion_stream_to_v2(_stream_from(items), request)
+    return [frame async for frame in gen]
+
+
+# ======================================================================
+# Low-level helpers: _sse, _emit, _DONE_FRAME
+# ======================================================================
+
+
+class TestSSEHelpers:
+    def test_sse_wraps_payload(self):
+        assert _sse("hello") == "data: hello\n\n"
+
+    def test_done_frame_constant(self):
+        # cohere-python and Fern-generated clients key off this exact
+        # sentinel; keep it byte-for-byte stable.
+        assert _DONE_FRAME == "data: [DONE]\n\n"
+
+    def test_emit_serializes_event_with_type_discriminator(self):
+        frame = _emit(
+            MessageStartEvent(id="abc", delta={"message": {"role": "assistant"}})
+        )
+        payload = _parse_event(frame)
+        assert payload["type"] == "message-start"
+        assert payload["id"] == "abc"
+        assert payload["delta"] == {"message": {"role": "assistant"}}
+
+
+# ======================================================================
+# _StreamState
+# ======================================================================
+
+
+class TestStreamState:
+    def test_content_index_monotonic(self):
+        st = _StreamState()
+        assert st.next_content_index() == 0
+        assert st.next_content_index() == 1
+        assert st.next_content_index() == 2
+
+    def test_citation_index_separate_from_content_index(self):
+        st = _StreamState()
+        st.next_content_index()  # 0
+        st.next_content_index()  # 1
+        # Citation indexing is independent of content indexing.
+        assert st.next_citation_index() == 0
+        assert st.next_citation_index() == 1
+
+
+# ======================================================================
+# _close_open_blocks
+# ======================================================================
+
+
+class TestCloseOpenBlocks:
+    def test_no_block_open_emits_nothing(self):
+        serving = _serving()
+        state = _StreamState()
+        assert serving._close_open_blocks(state) == []
+
+    def test_text_block_emits_content_end(self):
+        serving = _serving()
+        state = _StreamState()
+        state.active_block = ContentBlockType.TEXT
+        state.active_block_index = 1
+        out = serving._close_open_blocks(state)
+        assert len(out) == 1
+        payload = _parse_event(out[0])
+        assert payload == {"type": "content-end", "index": 1}
+        assert state.active_block is None
+        assert state.active_block_index is None
+
+    def test_thinking_block_emits_content_end(self):
+        serving = _serving()
+        state = _StreamState()
+        state.active_block = ContentBlockType.THINKING
+        state.active_block_index = 3
+        out = serving._close_open_blocks(state)
+        payload = _parse_event(out[0])
+        assert payload == {"type": "content-end", "index": 3}
+
+    def test_tool_call_block_emits_tool_call_end(self):
+        serving = _serving()
+        state = _StreamState()
+        state.active_block = ContentBlockType.TOOL_CALL
+        state.active_tool_index = 7
+        out = serving._close_open_blocks(state)
+        payload = _parse_event(out[0])
+        assert payload == {"type": "tool-call-end", "index": 7}
+        assert state.active_tool_index is None
+
+
+# ======================================================================
+# _handle_text_delta
+# ======================================================================
+
+
+class TestHandleTextDelta:
+    def test_opens_block_first_time(self):
+        serving = _serving()
+        state = _StreamState()
+        events = serving._handle_text_delta(state, "Hi")
+        assert len(events) == 2
+        start = _parse_event(events[0])
+        delta = _parse_event(events[1])
+        assert start["type"] == "content-start"
+        assert start["index"] == 0
+        assert start["delta"]["message"]["content"]["type"] == "text"
+        assert delta["type"] == "content-delta"
+        assert delta["index"] == 0
+        assert delta["delta"]["message"]["content"]["text"] == "Hi"
+        assert state.active_block == ContentBlockType.TEXT
+        assert state.active_block_index == 0
+
+    def test_continues_block_with_just_delta(self):
+        serving = _serving()
+        state = _StreamState()
+        serving._handle_text_delta(state, "Hi")
+        events = serving._handle_text_delta(state, " there")
+        # Only a delta event, no new content-start.
+        assert len(events) == 1
+        delta = _parse_event(events[0])
+        assert delta["type"] == "content-delta"
+        assert delta["delta"]["message"]["content"]["text"] == " there"
+
+    def test_switches_from_thinking_block(self):
+        serving = _serving(is_reasoning_model=True)
+        state = _StreamState()
+        # Open a thinking block first, then switch to text.
+        serving._handle_thinking_delta(state, "ponder")
+        events = serving._handle_text_delta(state, "answer")
+        types = [_parse_event(ev)["type"] for ev in events]
+        assert types == ["content-end", "content-start", "content-delta"]
+        # The text block gets a new index (1), distinct from thinking's 0.
+        assert _parse_event(events[1])["index"] == 1
+        assert state.active_block == ContentBlockType.TEXT
+
+
+# ======================================================================
+# _handle_thinking_delta
+# ======================================================================
+
+
+class TestHandleThinkingDelta:
+    def test_reasoning_model_opens_thinking_block(self):
+        serving = _serving(is_reasoning_model=True)
+        state = _StreamState()
+        events = serving._handle_thinking_delta(state, "thought")
+        assert len(events) == 2
+        start = _parse_event(events[0])
+        delta = _parse_event(events[1])
+        assert start["type"] == "content-start"
+        assert start["delta"]["message"]["content"]["type"] == "thinking"
+        assert delta["type"] == "content-delta"
+        assert delta["delta"]["message"]["content"]["thinking"] == "thought"
+        assert state.active_block == ContentBlockType.THINKING
+
+    def test_reasoning_model_continues_thinking_block(self):
+        serving = _serving(is_reasoning_model=True)
+        state = _StreamState()
+        serving._handle_thinking_delta(state, "first")
+        events = serving._handle_thinking_delta(state, " more")
+        assert len(events) == 1
+        assert _parse_event(events[0])["type"] == "content-delta"
+
+    def test_non_reasoning_model_emits_tool_plan_delta(self):
+        # Older Command models stream reasoning as ``tool_plan`` deltas;
+        # no content-start/end pair is emitted.
+        serving = _serving(is_reasoning_model=False)
+        state = _StreamState()
+        events = serving._handle_thinking_delta(state, "planning")
+        assert len(events) == 1
+        payload = _parse_event(events[0])
+        assert payload["type"] == "tool-plan-delta"
+        assert payload["delta"]["message"]["tool_plan"] == "planning"
+        # ``tool_plan`` deltas don't claim an active content block.
+        assert state.active_block is None
+
+    def test_non_reasoning_model_closes_open_text_block(self):
+        serving = _serving(is_reasoning_model=False)
+        state = _StreamState()
+        # Open a text block first.
+        serving._handle_text_delta(state, "answer")
+        events = serving._handle_thinking_delta(state, "rethink")
+        types = [_parse_event(ev)["type"] for ev in events]
+        assert types == ["content-end", "tool-plan-delta"]
+
+
+# ======================================================================
+# _handle_tool_call_deltas
+# ======================================================================
+
+
+class TestHandleToolCallDeltas:
+    def test_new_tool_call_opens_tool_call_start(self):
+        serving = _serving()
+        state = _StreamState()
+        deltas = [
+            type(
+                "Delta",
+                (),
+                {
+                    "index": 0,
+                    "id": "c1",
+                    "function": type(
+                        "Fn", (), {"name": "calc", "arguments": '{"x":'}
+                    )(),
+                },
+            )()
+        ]
+        events = serving._handle_tool_call_deltas(state, deltas)
+        assert len(events) == 1
+        payload = _parse_event(events[0])
+        assert payload["type"] == "tool-call-start"
+        assert payload["index"] == 0
+        tc = payload["delta"]["message"]["tool_calls"]
+        assert tc["id"] == "c1"
+        assert tc["function"]["name"] == "calc"
+        assert tc["function"]["arguments"] == '{"x":'
+        assert state.active_block == ContentBlockType.TOOL_CALL
+        assert state.active_tool_index == 0
+        assert 0 in state.tool_calls_seen
+
+    def test_subsequent_arguments_emit_delta(self):
+        serving = _serving()
+        state = _StreamState()
+        # First call: start.
+        first = [
+            type(
+                "Delta",
+                (),
+                {
+                    "index": 0,
+                    "id": "c1",
+                    "function": type("Fn", (), {"name": "calc", "arguments": ""})(),
+                },
+            )()
+        ]
+        serving._handle_tool_call_deltas(state, first)
+        # Second call: same index, additional arguments fragment.
+        more = [
+            type(
+                "Delta",
+                (),
+                {
+                    "index": 0,
+                    "id": None,
+                    "function": type("Fn", (), {"name": None, "arguments": "1}"})(),
+                },
+            )()
+        ]
+        events = serving._handle_tool_call_deltas(state, more)
+        assert len(events) == 1
+        payload = _parse_event(events[0])
+        assert payload["type"] == "tool-call-delta"
+        assert payload["index"] == 0
+        assert (
+            payload["delta"]["message"]["tool_calls"]["function"]["arguments"] == "1}"
+        )
+
+    def test_new_tool_call_closes_existing_content_block(self):
+        serving = _serving()
+        state = _StreamState()
+        # Open a text block, then start a tool call.
+        serving._handle_text_delta(state, "I'll call:")
+        deltas = [
+            type(
+                "Delta",
+                (),
+                {
+                    "index": 0,
+                    "id": "c1",
+                    "function": type("Fn", (), {"name": "calc", "arguments": "{}"})(),
+                },
+            )()
+        ]
+        events = serving._handle_tool_call_deltas(state, deltas)
+        types = [_parse_event(ev)["type"] for ev in events]
+        assert types == ["content-end", "tool-call-start"]
+
+
+# ======================================================================
+# _handle_citation_deltas
+# ======================================================================
+
+
+class TestHandleCitationDeltas:
+    """Test the per-delta emitter. Sources are pre-resolved by the
+    reasoning parser via the position map forwarded through
+    ``chat_template_kwargs`` (see :func:`_melody_sources_to_vllm` in
+    ``vllm/reasoning/cohere_command_reasoning_parser.py``), so the
+    handler only has to coerce to the SDK wire shape, apply the
+    ``THINKING_CONTENT`` -> ``PLAN`` rewrite, and drop citations with
+    no attributable sources.
+    """
+
+    @staticmethod
+    def _resolved_source(**overrides: Any) -> CitationSource:
+        base = dict(
+            type="document",
+            id="doc_hamlet",
+            document={"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
+        )
+        base.update(overrides)
+        return CitationSource(**base)
+
+    def test_resolved_citation_emits_start_and_end_with_real_id(self):
+        serving = _serving()
+        state = _StreamState()
+        citations = [
+            VLLMCitation(
+                start=0,
+                end=11,
+                text="Shakespeare",
+                sources=[self._resolved_source()],
+                type="TEXT_CONTENT",
+            )
+        ]
+        events = serving._handle_citation_deltas(state, citations)
+        assert len(events) == 2
+        start = _parse_event(events[0])
+        end = _parse_event(events[1])
+        assert start["type"] == "citation-start"
+        assert start["index"] == 0
+        cit_payload = start["delta"]["message"]["citations"]
+        assert cit_payload["start"] == 0
+        assert cit_payload["end"] == 11
+        assert cit_payload["text"] == "Shakespeare"
+        assert cit_payload["sources"] == [
+            {
+                "type": "document",
+                "id": "doc_hamlet",
+                "document": {"id": "doc_hamlet", "text": "Hamlet by Shakespeare."},
+            }
+        ]
+        assert end == {"type": "citation-end", "index": 0}
+
+    def test_multiple_sources_survive_to_wire(self):
+        serving = _serving()
+        state = _StreamState()
+        citations = [
+            VLLMCitation(
+                start=0,
+                end=5,
+                text="works",
+                sources=[
+                    self._resolved_source(id="d0", document={"id": "d0"}),
+                    self._resolved_source(id="d1", document={"id": "d1"}),
+                ],
+                type="TEXT_CONTENT",
+            )
+        ]
+        events = serving._handle_citation_deltas(state, citations)
+        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
+        ids = [s["id"] for s in cit_payload["sources"]]
+        assert ids == ["d0", "d1"]
+
+    def test_source_without_id_streamed_with_payload(self):
+        # Match the cohere api: id-less tool sources (from text-only
+        # tool results) still ride through with their ``tool_output``
+        # payload; only the ``id`` field is omitted.
+        serving = _serving()
+        state = _StreamState()
+        citations = [
+            VLLMCitation(
+                start=0,
+                end=5,
+                text="hello",
+                sources=[CitationSource(type="tool", tool_output={"content": "hi"})],
+            )
+        ]
+        events = serving._handle_citation_deltas(state, citations)
+        assert len(events) >= 1
+        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
+        assert cit_payload["sources"] == [
+            {"type": "tool", "tool_output": {"content": "hi"}}
+        ]
+        assert "id" not in cit_payload["sources"][0]
+
+    def test_citation_with_no_sources_streamed_dropped(self):
+        # Citations whose source list is entirely empty are still
+        # dropped so we don't emit an unattributed anchor.
+        serving = _serving()
+        state = _StreamState()
+        citations = [
+            VLLMCitation(
+                start=0,
+                end=5,
+                text="hello",
+                sources=[],
+            )
+        ]
+        events = serving._handle_citation_deltas(state, citations)
+        assert events == []
+
+    def test_plan_rewrite_on_non_reasoning_model(self):
+        serving = _serving(is_reasoning_model=False)
+        state = _StreamState()
+        citations = [
+            VLLMCitation(
+                start=0,
+                end=5,
+                text="hello",
+                sources=[self._resolved_source()],
+                type="THINKING_CONTENT",
+            )
+        ]
+        events = serving._handle_citation_deltas(state, citations)
+        cit_payload = _parse_event(events[0])["delta"]["message"]["citations"]
+        assert cit_payload["type"] == "PLAN"
+
+
+# ======================================================================
+# _build_message_end_event
+# ======================================================================
+
+
+class TestBuildMessageEndEvent:
+    def test_without_usage(self):
+        serving = _serving()
+        frame = serving._build_message_end_event(chunk_id="abc", finish_reason="stop")
+        payload = _parse_event(frame)
+        assert payload["type"] == "message-end"
+        assert payload["id"] == "abc"
+        assert payload["delta"]["finish_reason"] == "COMPLETE"
+        assert "usage" not in payload["delta"]
+
+    def test_with_usage(self):
+        serving = _serving()
+        chunk = ChatCompletionStreamResponse.model_validate(
+            {
+                "id": "abc",
+                "object": "chat.completion.chunk",
+                "created": 0,
+                "model": "m",
+                "choices": [],
+                "usage": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 5,
+                    "total_tokens": 15,
+                },
+            }
+        )
+        frame = serving._build_message_end_event(
+            chunk_id="abc", finish_reason="length", usage_chunk=chunk
+        )
+        payload = _parse_event(frame)
+        assert payload["delta"]["finish_reason"] == "MAX_TOKENS"
+        usage = payload["delta"]["usage"]
+        assert usage["billed_units"] == {"input_tokens": 10, "output_tokens": 5}
+        assert usage["tokens"] == {"input_tokens": 10, "output_tokens": 5}
+        assert "cached_tokens" not in usage
+
+    def test_with_cached_tokens(self):
+        serving = _serving()
+        chunk = ChatCompletionStreamResponse.model_validate(
+            {
+                "id": "abc",
+                "object": "chat.completion.chunk",
+                "created": 0,
+                "model": "m",
+                "choices": [],
+                "usage": {
+                    "prompt_tokens": 10,
+                    "completion_tokens": 5,
+                    "total_tokens": 15,
+                    "prompt_tokens_details": {"cached_tokens": 3},
+                },
+            }
+        )
+        frame = serving._build_message_end_event(
+            chunk_id="abc", finish_reason="stop", usage_chunk=chunk
+        )
+        payload = _parse_event(frame)
+        assert payload["delta"]["usage"]["cached_tokens"] == 3
+
+
+# ======================================================================
+# End-to-end: _chat_completion_stream_to_v2
+# ======================================================================
+
+
+class TestChatCompletionStreamToV2:
+    """End-to-end stream lifecycle tests."""
+
+    @pytest.mark.asyncio
+    async def test_text_only_happy_path_emits_full_lifecycle(self):
+        serving = _serving(is_reasoning_model=True)
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(content="Hi"),
+            _make_chunk(content=" there"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 4,
+                    "completion_tokens": 2,
+                    "total_tokens": 6,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        # message-start, content-start, content-delta, content-delta,
+        # content-end, message-end, then [DONE] as the last frame.
+        assert types == [
+            "message-start",
+            "content-start",
+            "content-delta",
+            "content-delta",
+            "content-end",
+            "message-end",
+        ]
+        assert frames[-1] == _DONE_FRAME
+        # message-end should carry usage stats from the trailing chunk.
+        end_payload = _parse_event(frames[-2])
+        assert end_payload["delta"]["finish_reason"] == "COMPLETE"
+        assert end_payload["delta"]["usage"]["billed_units"]["input_tokens"] == 4
+
+    @pytest.mark.asyncio
+    async def test_text_then_tool_call_closes_text_first(self):
+        serving = _serving(is_reasoning_model=True)
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(content="planning..."),
+            _make_chunk(
+                tool_calls=[
+                    {
+                        "index": 0,
+                        "id": "c1",
+                        "type": "function",
+                        "function": {"name": "calc", "arguments": "{}"},
+                    }
+                ]
+            ),
+            _make_chunk(finish_reason="tool_calls"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        # Text block is opened+delta, then closed before the tool-call
+        # opens, and the final close happens on the usage chunk path.
+        assert types == [
+            "message-start",
+            "content-start",
+            "content-delta",
+            "content-end",
+            "tool-call-start",
+            "tool-call-end",
+            "message-end",
+        ]
+        # finish_reason captured from the prior chunk.
+        end = _parse_event(frames[-2])
+        assert end["delta"]["finish_reason"] == "TOOL_CALL"
+
+    @pytest.mark.asyncio
+    async def test_thinking_then_text_reasoning_model(self):
+        serving = _serving(is_reasoning_model=True)
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(reasoning="thinking..."),
+            _make_chunk(content="answer"),
+            _make_chunk(finish_reason="stop"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        # Thinking block opens with index 0; text block reopens with index 1.
+        assert types == [
+            "message-start",
+            "content-start",
+            "content-delta",
+            "content-end",
+            "content-start",
+            "content-delta",
+            "content-end",
+            "message-end",
+        ]
+        thinking_start = _parse_event(frames[1])
+        assert thinking_start["delta"]["message"]["content"]["type"] == "thinking"
+        text_start = _parse_event(frames[4])
+        assert text_start["delta"]["message"]["content"]["type"] == "text"
+
+    @pytest.mark.asyncio
+    async def test_reasoning_on_non_reasoning_model_emits_tool_plan_delta(self):
+        serving = _serving(is_reasoning_model=False)
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(reasoning="planning"),
+            _make_chunk(
+                tool_calls=[
+                    {
+                        "index": 0,
+                        "id": "c1",
+                        "type": "function",
+                        "function": {"name": "f", "arguments": "{}"},
+                    }
+                ]
+            ),
+            _make_chunk(finish_reason="tool_calls"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        assert types == [
+            "message-start",
+            "tool-plan-delta",
+            "tool-call-start",
+            "tool-call-end",
+            "message-end",
+        ]
+        # No thinking content blocks should be present.
+        assert "content-start" not in types
+        assert "content-end" not in types
+
+    @pytest.mark.asyncio
+    async def test_done_marker_in_middle_closes_open_block(self):
+        # Some upstreams send [DONE] without a trailing usage-only chunk.
+        # The translator must still emit message-end before [DONE] so
+        # Cohere clients don't hang.
+        serving = _serving(is_reasoning_model=True)
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(content="Hi"),
+            _make_done(),
+        ]
+        frames = await _drain(serving, items)
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        assert types == [
+            "message-start",
+            "content-start",
+            "content-delta",
+            "content-end",
+            "message-end",
+        ]
+        assert frames[-1] == _DONE_FRAME
+
+    @pytest.mark.asyncio
+    async def test_skips_empty_and_non_data_lines(self):
+        serving = _serving()
+        items = [
+            "\n",
+            "event: ping\n\n",
+            _make_chunk(role="assistant"),
+            "data: \n\n",  # empty data
+            _make_chunk(content="Hi"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        assert frames[-1] == _DONE_FRAME
+        # Should still produce a complete lifecycle.
+        types = [_parse_event(f)["type"] for f in frames[:-1]]
+        assert types[0] == "message-start"
+        assert types[-1] == "message-end"
+
+    @pytest.mark.asyncio
+    async def test_exception_in_chunk_parsing_emits_error_message_end(self):
+        # An invalid JSON payload after the first chunk triggers the
+        # error path: a synthetic message-end with finish_reason=ERROR
+        # followed by [DONE].
+        serving = _serving()
+        items = [
+            _make_chunk(role="assistant"),
+            "data: {not valid json}\n\n",
+        ]
+        frames = await _drain(serving, items)
+        assert frames[-1] == _DONE_FRAME
+        # Find the error-shaped message-end.
+        error_end = _parse_event(frames[-2])
+        assert error_end["type"] == "message-end"
+        assert error_end["delta"]["finish_reason"] == "ERROR"
+        assert "error" in error_end["delta"]
+
+    @pytest.mark.asyncio
+    async def test_error_message_end_is_sanitized(self):
+        # The client-visible ``error`` field on the terminal message-end
+        # must be routed through ``sanitize_message`` so tracebacks,
+        # host filesystem paths, and Python object memory addresses
+        # can't leak from a mid-stream exception into the SSE payload.
+        serving = _serving()
+
+        async def _raises_mid_stream() -> AsyncGenerator[str, None]:
+            yield _make_chunk(role="assistant")
+            raise RuntimeError(
+                "boom <Foo at 0x7fabc0deface> reading /Users/dev/vllm/x.py"
+            )
+
+        request = CohereChatV2Request(
+            model="m",
+            messages=[{"role": "user", "content": "hi"}],
+            stream=True,
+        )
+        frames = [
+            f
+            async for f in serving._chat_completion_stream_to_v2(
+                _raises_mid_stream(), request
+            )
+        ]
+        assert frames[-1] == _DONE_FRAME
+        error_end = _parse_event(frames[-2])
+        assert error_end["delta"]["finish_reason"] == "ERROR"
+        err_msg = error_end["delta"]["error"]
+        # Memory address stripped and path replaced with ``<path>``.
+        assert "0x7fabc0deface" not in err_msg
+        assert "/Users/dev/vllm/x.py" not in err_msg
+        assert "<path>" in err_msg
+        # The human-facing prefix should still survive sanitization.
+        assert "boom" in err_msg
+
+    @pytest.mark.asyncio
+    async def test_citations_in_delta_emit_citation_events(self):
+        # End-to-end streaming citation flow. The parser hands us
+        # already-resolved sources (see ``_melody_sources_to_vllm``
+        # in ``vllm/reasoning/cohere_command_reasoning_parser.py``);
+        # the streaming loop just has to coerce and emit.
+        serving = _serving()
+        items = [
+            _make_chunk(role="assistant"),
+            _make_chunk(content="hello"),
+            _make_chunk(
+                citations=[
+                    {
+                        "start": 0,
+                        "end": 5,
+                        "text": "hello",
+                        "sources": [
+                            {
+                                "type": "document",
+                                "id": "d1",
+                                "document": {"id": "d1", "text": "cited"},
+                            }
+                        ],
+                        "type": "TEXT_CONTENT",
+                    }
+                ]
+            ),
+            _make_chunk(finish_reason="stop"),
+            _make_chunk(
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        parsed = [_parse_event(f) for f in frames[:-1]]
+        types = [p["type"] for p in parsed]
+        assert "citation-start" in types
+        assert "citation-end" in types
+        cit_start = next(p for p in parsed if p["type"] == "citation-start")
+        assert cit_start["delta"]["message"]["citations"]["sources"] == [
+            {
+                "type": "document",
+                "id": "d1",
+                "document": {"id": "d1", "text": "cited"},
+            }
+        ]
+
+    @pytest.mark.asyncio
+    async def test_first_chunk_emits_message_start_with_chunk_id(self):
+        serving = _serving()
+        items = [
+            _make_chunk(chunk_id="my-id", role="assistant"),
+            _make_chunk(chunk_id="my-id", content="hi"),
+            _make_chunk(
+                chunk_id="my-id",
+                omit_choices=True,
+                usage={
+                    "prompt_tokens": 1,
+                    "completion_tokens": 1,
+                    "total_tokens": 2,
+                },
+            ),
+        ]
+        frames = await _drain(serving, items)
+        ms = _parse_event(frames[0])
+        assert ms["type"] == "message-start"
+        assert ms["id"] == "my-id"
+        assert ms["delta"]["message"]["role"] == "assistant"

tests/reasoning/test_cohere_command_reasoning_parser.py

@@ -25,6 +25,8 @@ from vllm.reasoning.cohere_command_reasoning_parser import (
     CohereCommand3ReasoningParser,
     CohereCommand4ReasoningParser,
     _has_effective_tools,
+    _melody_citations_to_vllm,
+    _melody_sources_to_vllm,
     _response_format_type,
     _schema_dict_from_structured_outputs,
     convert_schema_to_structural_tags,
@@ -623,3 +625,68 @@ class TestAdjustRequestTools:
         types = _content_types(o.structured_outputs.structural_tag)
         assert "grammar" in types
         assert "json_schema" in types
+
+
+class TestMelodySourceResolution:
+    """Pin the parser-side source resolution (see
+    ``_melody_sources_to_vllm``). The parser receives melody's numeric
+    ``(tool_call_index, tool_result_indices)`` addressing and resolves
+    it against a ``position_to_source`` map handed in by
+    ``CohereServingChatV2._apply_cohere_template_kwargs`` -- the same
+    map that would otherwise live in the serving layer's resolver.
+    """
+
+    @staticmethod
+    def _fake_melody_source(bucket: int, indices: list[int]) -> Any:
+        return SimpleNamespace(tool_call_index=bucket, tool_result_indices=indices)
+
+    def test_multi_index_source_fans_out(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0", document={"id": "d0"}),
+            (0, 1): CitationSource(type="document", id="d1", document={"id": "d1"}),
+        }
+        raw = [self._fake_melody_source(0, [0, 1])]
+        out = _melody_sources_to_vllm(raw, position_map)
+        assert [s.id for s in out] == ["d0", "d1"]
+        # Verify type / payload were plumbed through, not just the id.
+        assert out[0].type == "document"
+        assert out[0].document == {"id": "d0"}
+
+    def test_unresolvable_position_skipped(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0"),
+        }
+        raw = [self._fake_melody_source(9, [0])]
+        assert _melody_sources_to_vllm(raw, position_map) == []
+
+    def test_missing_position_map_drops_all_sources(self):
+        # A parser instance without a position map (parser wired
+        # outside of ``CohereServingChatV2``) can't attribute anything,
+        # so every source is dropped. Callers downstream will see the
+        # citation with empty ``sources`` and drop it entirely.
+        raw = [self._fake_melody_source(0, [0])]
+        assert _melody_sources_to_vllm(raw, None) == []
+
+    def test_citations_pass_through_is_thinking_tag(self):
+        from vllm.entrypoints.cohere.cohere_chat_message import CitationSource
+
+        position_map: dict[tuple[int, int], CitationSource] = {
+            (0, 0): CitationSource(type="document", id="d0"),
+        }
+        raw = [
+            SimpleNamespace(
+                start_index=0,
+                end_index=5,
+                text="hello",
+                is_thinking=True,
+                sources=[self._fake_melody_source(0, [0])],
+            )
+        ]
+        out = _melody_citations_to_vllm(raw, position_map)
+        assert out is not None
+        assert out[0].type == "THINKING_CONTENT"
+        assert out[0].sources[0].id == "d0"

vllm/config/model.py

@@ -93,6 +93,7 @@ TokenizerMode = Literal[
     "deepseek_v4",
     "inkling",
     "kimi_k3",
+    "cohere",
 ]
 ModelDType = Literal["auto", "half", "float16", "bfloat16", "float", "float32"]
 LogprobsMode = Literal[
@@ -151,6 +152,9 @@ class ModelConfig:
     - "deepseek_v4" will always use the tokenizer from `deepseek_v4`.
     - "kimi_k3" will always use the "hf" tokenizer but render chat prompts
       with Kimi K3's Python XTML encoding instead of a Jinja template.
+    - "cohere" uses the standard HF tokenizer but renders the chat template
+      via the `cohere_melody` library (cmd3 / cmd4 templates) instead of
+      Jinja, and surfaces grounded-citation metadata on responses.
     - Other custom values can be supported via plugins.
 
     To swap the Rust BPE backend that powers HF fast tokenizers for the

vllm/entrypoints/cohere/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

vllm/entrypoints/cohere/api_router.py

@@ -0,0 +1,244 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""FastAPI router for the Cohere Chat v2 API (``POST /cohere/v2/chat``).
+
+The Cohere v2 protocol models are sourced from the official ``cohere``
+Python SDK (``pip install cohere``). To keep that an *optional*
+dependency for vLLM, the SDK-dependent imports - and the route handler
+itself - are gated on a one-shot probe at module load. If the SDK isn't
+installed, :func:`attach_router` becomes a no-op (with an info log) and
+vLLM continues to boot normally.
+
+Even when the SDK is installed, :func:`attach_router` also requires
+``VLLM_ENABLE_COHERE_API=1`` in the environment before it will expose
+the route. This keeps non-Cohere deployments that pull in the SDK for
+unrelated reasons (e.g. test dependencies) from accidentally exposing
+the api.
+
+Note: the handler must live at module scope (not inside
+``attach_router``) so that FastAPI's ``typing.get_type_hints`` resolves
+the ``CohereChatV2Request`` body annotation against the module's
+globals. Defining it locally inside ``attach_router`` would hide the
+type from ``get_type_hints``, causing FastAPI to silently degrade the
+body parameter into a query parameter and reject every request with
+422.
+"""
+
+import json
+from http import HTTPStatus
+
+from fastapi import APIRouter, Depends, FastAPI, Request
+from fastapi.responses import JSONResponse, Response, StreamingResponse
+from starlette.middleware.base import BaseHTTPMiddleware
+
+import vllm.envs as envs
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse
+from vllm.entrypoints.serve.utils.api_utils import (
+    load_aware_call,
+    sanitize_message,
+    validate_json_request,
+    with_cancellation,
+)
+from vllm.logger import init_logger
+
+_COHERE_PATH_PREFIX = "/cohere/"
+
+logger = init_logger(__name__)
+
+
+try:
+    import cohere  # noqa: F401  -- dependency probe
+except ImportError:
+    _SDK_AVAILABLE = False
+else:
+    _SDK_AVAILABLE = True
+
+
+if _SDK_AVAILABLE:
+    from vllm.entrypoints.cohere.protocol import (
+        CohereChatV2Request,
+        CohereChatV2Response,
+        CohereError,
+    )
+    from vllm.entrypoints.cohere.serving import CohereServingChatV2
+
+    router = APIRouter()
+
+    def _serving(request: Request) -> CohereServingChatV2 | None:
+        return getattr(request.app.state, "cohere_serving_chat_v2", None)
+
+    def _request_id(raw_request: Request | None) -> str | None:
+        """Best-effort lookup of the active request id.
+
+        Prefers the id the underlying chat handler stamped onto
+        ``raw_request.state.request_metadata`` (if it got that far before
+        failing), falling back to the ``X-Request-Id`` HTTP header. May
+        return ``None`` if neither is available, in which case the field
+        is omitted from the response.
+        """
+        if raw_request is None:
+            return None
+        meta = getattr(raw_request.state, "request_metadata", None)
+        if meta is not None and getattr(meta, "request_id", None):
+            return meta.request_id
+        return raw_request.headers.get("X-Request-Id")
+
+    def _error_response(
+        error: ErrorResponse,
+        raw_request: Request | None,
+        *,
+        fallback_status: int = HTTPStatus.BAD_REQUEST,
+    ) -> JSONResponse:
+        """Translate vLLM's internal error envelope into Cohere's shape."""
+        info = error.error
+        status = info.code or fallback_status
+        return JSONResponse(
+            status_code=status,
+            content=CohereError(
+                message=sanitize_message(info.message),
+                id=_request_id(raw_request),
+            ).model_dump(exclude_none=True),
+        )
+
+    @router.post(
+        "/cohere/v2/chat",
+        dependencies=[Depends(validate_json_request)],
+        responses={
+            HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
+            HTTPStatus.BAD_REQUEST.value: {"model": CohereError},
+            HTTPStatus.NOT_FOUND.value: {"model": CohereError},
+            HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": CohereError},
+        },
+    )
+    @with_cancellation
+    @load_aware_call
+    async def chat_v2(request: CohereChatV2Request, raw_request: Request):
+        handler = _serving(raw_request)
+        if handler is None:
+            return JSONResponse(
+                status_code=HTTPStatus.NOT_IMPLEMENTED.value,
+                content=CohereError(
+                    message="The model does not support the Cohere v2 chat API.",
+                    id=_request_id(raw_request),
+                ).model_dump(exclude_none=True),
+            )
+
+        try:
+            result = await handler.create_chat_v2(request, raw_request)
+        except Exception as e:  # noqa: BLE001 - report as 500 for parity
+            logger.exception("Error in /cohere/v2/chat: %s", e)
+            return JSONResponse(
+                status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value,
+                content=CohereError(
+                    message=sanitize_message(str(e)),
+                    id=_request_id(raw_request),
+                ).model_dump(exclude_none=True),
+            )
+
+        match result:
+            case ErrorResponse():
+                return _error_response(result, raw_request)
+            case CohereChatV2Response():
+                return JSONResponse(content=result.model_dump(exclude_none=True))
+            case _:
+                return StreamingResponse(content=result, media_type="text/event-stream")
+
+    class CohereErrorEnvelopeMiddleware(BaseHTTPMiddleware):
+        """Rewrite vLLM error bodies into the Cohere ``{message, id}`` shape.
+
+        The endpoint handler above already returns :class:`CohereError` for
+        errors it owns, but globally-registered exception handlers (e.g.
+        :func:`validation_exception_handler` for pydantic body errors,
+        :func:`http_exception_handler`, engine error handlers) fire *before*
+        the handler runs and produce vLLM's internal
+        ``ErrorResponse`` shape (``{"error": {"message": ...}}``). That
+        shape doesn't match the ``CohereError`` schema advertised on the
+        route's OpenAPI ``responses``, so clients (and schema-conformance
+        tests like ``test_openai_schema.py``) would see a mismatch on
+        those paths. This middleware normalises any error body on
+        ``/cohere/*`` responses to :class:`CohereError`.
+        """
+
+        async def dispatch(self, request: Request, call_next):
+            response = await call_next(request)
+            if not request.url.path.startswith(_COHERE_PATH_PREFIX):
+                return response
+            if response.status_code < 400:
+                return response
+            content_type = response.headers.get("content-type", "")
+            if not content_type.startswith("application/json"):
+                return response
+
+            body = b"".join([chunk async for chunk in response.body_iterator])
+            translated = _translate_vllm_error_body(body, request)
+            if translated is not None:
+                return translated
+            passthrough_headers = {
+                k: v
+                for k, v in response.headers.items()
+                if k.lower() != "content-length"
+            }
+            return Response(
+                content=body,
+                status_code=response.status_code,
+                headers=passthrough_headers,
+                media_type=content_type,
+            )
+
+    def _translate_vllm_error_body(raw: bytes, request: Request) -> JSONResponse | None:
+        """Translate a vLLM ``ErrorResponse`` body to a ``CohereError`` body.
+
+        Returns ``None`` if ``raw`` does not match the vLLM error envelope
+        (which signals the middleware to pass the body through unchanged).
+        """
+        try:
+            data = json.loads(raw)
+        except (json.JSONDecodeError, TypeError, ValueError):
+            return None
+        if not (
+            isinstance(data, dict)
+            and isinstance(data.get("error"), dict)
+            and "message" in data["error"]
+        ):
+            return None
+        try:
+            err = ErrorResponse.model_validate(data)
+        except Exception:  # noqa: BLE001 - malformed envelope; pass through
+            return None
+        return _error_response(err, request)
+
+
+def attach_router(app: FastAPI) -> None:
+    """Register ``POST /cohere/v2/chat`` on ``app``.
+
+    No-op when either:
+
+    * the ``VLLM_ENABLE_COHERE_API`` env var isn't set to ``1``. The
+      Cohere v2 endpoint is opt-in because it carries Cohere-specific
+      request/response semantics (grounding citations, tool_plan,
+      PLAN/THINKING_CONTENT blocks) that are only meaningful when
+      serving a Cohere Command-family model.
+    * the optional ``cohere`` SDK isn't installed (the v2 protocol
+      models live there)
+
+    The two skip paths log at different levels: an operator who set
+    ``VLLM_ENABLE_COHERE_API=1`` but forgot to install ``cohere`` sees
+    a WARNING (they explicitly asked for the endpoint and it's silently
+    absent), whereas the default-off skip logs at debug.
+    """
+    enabled = envs.VLLM_ENABLE_COHERE_API
+    if not enabled:
+        logger.debug(
+            "VLLM_ENABLE_COHERE_API is not set; /cohere/v2/chat endpoint "
+            "disabled. Set VLLM_ENABLE_COHERE_API=1 to enable it."
+        )
+        return
+    if not _SDK_AVAILABLE:
+        logger.warning(
+            "VLLM_ENABLE_COHERE_API=1 but the `cohere` SDK is not "
+            "installed; /cohere/v2/chat will not be exposed. Install "
+            "with `pip install cohere` to enable the endpoint."
+        )
+        return
+    app.include_router(router)
+    app.add_middleware(CohereErrorEnvelopeMiddleware)

vllm/entrypoints/cohere/cohere_chat_message.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vLLM chat-protocol extensions for grounded (citation-carrying) models.
+
+This module keeps the OpenAI chat completion protocol classes
+(:class:`ChatMessage` / :class:`DeltaMessage`) free of Cohere-specific fields.
+It provides two things:
+
+* :class:`Citation` / :class:`CitationSource`: the vLLM-internal citation
+  representation produced by the Cohere reasoning parser
+  (:mod:`vllm.reasoning.cohere_command_reasoning_parser`) and consumed by
+  :mod:`vllm.entrypoints.cohere.serving`. This is *not* the on-wire Cohere
+  SDK shape -- that conversion happens in the serving layer.
+* :class:`CohereChatMessage` / :class:`CohereDeltaMessage`: :class:`ChatMessage`
+  / :class:`DeltaMessage` subclasses that add a ``citations`` field. They
+  are only instantiated by the citation-aware serving handler
+  (:class:`vllm.entrypoints.cohere.serving.CohereServingChatV2`) and by the
+  Cohere reasoning parser's streaming path. Non-Cohere code paths continue
+  to construct plain :class:`ChatMessage` / :class:`DeltaMessage`.
+
+The module intentionally does *not* import the ``cohere`` Python SDK so the
+reasoning parser (which is Cohere-model-specific but ships without the SDK
+dependency) can import from here freely.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import Field, model_serializer
+
+from vllm.entrypoints.openai.chat_completion.protocol import ChatMessage
+from vllm.entrypoints.openai.engine.protocol import (
+    DeltaMessage,
+    OpenAIBaseModel,
+)
+
+
+class CitationSource(OpenAIBaseModel):
+    """Source attribution for a :class:`Citation`.
+
+    Mirrors the shape used by Cohere's Chat v2 API. ``type`` is the source
+    discriminator (``document`` or ``tool``); ``id`` is the citing
+    document/tool-output identifier; ``document`` and ``tool_output`` carry
+    the original payload that produced the citation.
+
+    Sources are fully resolved by the reasoning parser at emit time
+    (see :func:`_melody_sources_to_vllm` in
+    :mod:`vllm.reasoning.cohere_command_reasoning_parser`) using the
+    ``POSITION_TO_SOURCE_KEY`` map forwarded through
+    ``chat_template_kwargs`` by
+    :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._apply_cohere_template_kwargs`.
+    That means every instance that reaches the serving / wire layer
+    already has ``type`` / ``id`` populated; the serving layer only has
+    to coerce to :class:`cohere.types.Citation` and rewrite
+    ``THINKING_CONTENT`` -> ``PLAN`` for non-reasoning models.
+    """
+
+    type: Literal["document", "tool"] | None = None
+    id: str | None = None
+    document: dict[str, Any] | None = None
+    tool_output: dict[str, Any] | None = None
+
+
+class Citation(OpenAIBaseModel):
+    """A citation grounding a span of generated text in source material.
+
+    vLLM-internal representation used by the Cohere reasoning parser and
+    the Cohere v2 serving layer. This is not the on-wire Cohere SDK shape:
+    conversion to the SDK's ``cohere.types.Citation`` happens in
+    :mod:`vllm.entrypoints.cohere.serving`.
+    """
+
+    start: int | None = None
+    """Start character offset in the surrounding text content."""
+    end: int | None = None
+    """End character offset (exclusive) in the surrounding text content."""
+    text: str | None = None
+    """The cited text snippet."""
+    sources: list[CitationSource] = Field(default_factory=list)
+    """Source documents / tool outputs that ground this citation."""
+    content_index: int | None = None
+    """Index of the content block this citation refers to (when the message
+    has multiple content blocks)."""
+    type: Literal["TEXT_CONTENT", "THINKING_CONTENT", "PLAN"] | None = None
+    """Which kind of content block this citation grounds: the user-visible
+    text (``TEXT_CONTENT``), a thinking block (``THINKING_CONTENT``), or a
+    tool-plan block (``PLAN``). ``None`` means unspecified."""
+
+
+class CohereChatMessage(ChatMessage):
+    """:class:`ChatMessage` extension carrying grounding citations.
+
+    Only instantiated by :class:`CohereServingChatV2` (and any other
+    :class:`OpenAIServingChat` subclass that opts in by overriding
+    ``_create_chat_message``). Regular OpenAI-compatible handlers keep
+    emitting plain :class:`ChatMessage` so their response schema is
+    unchanged.
+
+    The response envelope declares ``message: SerializeAsAny[ChatMessage]``,
+    so pydantic serializes this subclass with its own schema (including
+    ``citations``) when it flows through
+    :class:`ChatCompletionResponseChoice`.
+    """
+
+    citations: list[Citation] | None = None
+
+    @model_serializer(mode="wrap")
+    def _serialize(self, handler):
+        # ``mode="wrap"`` fully overrides (rather than chains) the
+        # parent's ``@model_serializer``, so we explicitly delegate via
+        # ``super()._serialize(handler)`` to preserve the parent's
+        # cleanup (e.g. stripping empty ``tool_calls``). Then we drop an
+        # unset ``citations`` field so the wire matches the OpenAI-style
+        # contract that optional vLLM extensions are omitted rather than
+        # serialized as ``null``.
+        data = super()._serialize(handler)
+        if not data.get("citations"):
+            data.pop("citations", None)
+        return data
+
+
+class CohereDeltaMessage(DeltaMessage):
+    """:class:`DeltaMessage` extension carrying grounding citations for streaming.
+
+    Emitted by
+    :class:`vllm.reasoning.cohere_command_reasoning_parser.BaseCohereCommandReasoningParser`
+    on delta events whose payload includes citations. Non-Cohere parsers
+    return plain :class:`DeltaMessage`, so their streamed shape is
+    unchanged.
+
+    The response envelope declares
+    ``delta: SerializeAsAny[DeltaMessage]``, so pydantic serializes this
+    subclass with its own schema when it flows through
+    :class:`ChatCompletionResponseStreamChoice`.
+    """
+
+    citations: list[Citation] | None = None
+
+    @model_serializer(mode="wrap")
+    def _serialize(self, handler):
+        # See ``CohereChatMessage._serialize`` for why we delegate via
+        # ``super()`` instead of calling ``handler(self)`` directly.
+        data = super()._serialize(handler)
+        if not data.get("citations"):
+            data.pop("citations", None)
+        return data

vllm/entrypoints/cohere/protocol.py

@@ -0,0 +1,362 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cohere Chat v2 API protocol.
+
+The bulk of the wire types come straight from the official ``cohere``
+Python SDK so we stay in sync with the upstream specification and
+avoid re-declaring the schema. We only own three things locally:
+
+1. The top-level request body model (the SDK doesn't ship one — its
+   ``ClientV2.chat`` takes the body as kwargs), with vLLM-specific
+   extensions (``kv_transfer_params`` / ``chat_template_kwargs``).
+2. The non-streaming response envelope (the SDK exposes the message
+   shape via :class:`AssistantMessageResponse` but no full response
+   wrapper).
+3. The streaming discriminated union (the SDK exports each event type
+   individually but not as a combined ``Annotated[Union[...],
+   discriminator]``).
+
+Importing this module pulls in the ``cohere`` package. The router that
+mounts ``POST /cohere/v2/chat`` guards on that import succeeding so vLLM still
+boots without the SDK installed.
+
+See https://docs.cohere.com/reference/chat for the upstream spec.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from cohere import types as _sdk
+from cohere.types import (
+    AssistantChatMessageV2,
+    AssistantMessageResponse,
+    ChatMessageV2,
+    ChatRequestSafetyMode,
+    Citation,
+    CitationOptions,
+    Document,
+    ResponseFormatV2,
+    SystemChatMessageV2,
+    Thinking,
+    ToolCallV2,
+    ToolChatMessageV2,
+    ToolV2,
+    UserChatMessageV2,
+)
+from pydantic import BaseModel, Field, field_validator
+
+# Re-export the SDK wire-format types alongside our local extensions so
+# ``vllm.entrypoints.cohere.serving`` and friends can import everything
+# they need from this module.
+__all__ = [
+    "AssistantChatMessageV2",
+    "AssistantMessageResponse",
+    "ChatMessageV2",
+    "Citation",
+    "CitationEndEvent",
+    "CitationOptions",
+    "CitationStartEvent",
+    "CohereChatV2Request",
+    "CohereChatV2Response",
+    "CohereError",
+    "CohereFinishReason",
+    "CohereLogprobItem",
+    "CohereUsage",
+    "CohereUsageBilledUnits",
+    "CohereUsageTokens",
+    "ContentDeltaEvent",
+    "ContentEndEvent",
+    "ContentStartEvent",
+    "Document",
+    "MessageEndEvent",
+    "MessageStartEvent",
+    "ResponseFormatV2",
+    "SystemChatMessageV2",
+    "Thinking",
+    "ToolCallDeltaEvent",
+    "ToolCallEndEvent",
+    "ToolCallStartEvent",
+    "ToolCallV2",
+    "ToolChatMessageV2",
+    "ToolPlanDeltaEvent",
+    "ToolV2",
+    "UserChatMessageV2",
+]
+
+
+# ---------------------------------------------------------------------------
+# Errors
+# ---------------------------------------------------------------------------
+
+
+class CohereError(BaseModel):
+    """Top-level error body returned by ``/cohere/v2/chat`` error responses.
+
+    Cohere's documented error schemas are uniform: ``{message, id}``.
+    """
+
+    message: str
+    id: str | None = None
+
+
+# ---------------------------------------------------------------------------
+# Tool choice / finish reasons
+# ---------------------------------------------------------------------------
+#
+# These literals aren't first-class enums in the SDK but are documented at
+# https://docs.cohere.com/reference/chat. We declare them here so the
+# request/response models can validate them.
+
+CohereToolChoice = Literal["REQUIRED", "NONE"]
+
+CohereFinishReason = Literal[
+    "COMPLETE",
+    "STOP_SEQUENCE",
+    "MAX_TOKENS",
+    "TOOL_CALL",
+    "ERROR",
+    "TIMEOUT",
+]
+
+
+# ---------------------------------------------------------------------------
+# Request
+# ---------------------------------------------------------------------------
+
+
+class CohereChatV2Request(BaseModel):
+    """Cohere Chat v2 request body.
+
+    Mirrors the schema documented at https://docs.cohere.com/reference/chat.
+    All structured fields delegate to the official SDK types so the body
+    schema stays in sync with the upstream spec.
+    """
+
+    model: str
+    messages: list[ChatMessageV2]
+    stream: bool | None = False
+
+    # Tooling
+    tools: list[ToolV2] | None = None
+    strict_tools: bool | None = None
+    tool_choice: CohereToolChoice | None = None
+
+    # Grounding
+    documents: list[str | Document] | None = None
+    citation_options: CitationOptions | None = None
+
+    # Output
+    response_format: ResponseFormatV2 | None = None
+    safety_mode: ChatRequestSafetyMode | None = None
+    max_tokens: int | None = None
+    stop_sequences: list[str] | None = None
+
+    # Sampling
+    temperature: float | None = None
+    seed: int | None = None
+    frequency_penalty: float | None = None
+    presence_penalty: float | None = None
+    k: int | None = None
+    p: float | None = None
+    logprobs: bool | None = None
+
+    # Reasoning
+    thinking: Thinking | None = None
+
+    # Scheduling
+    priority: int | None = None
+
+    # vLLM-specific extensions (not in Cohere spec). These mirror what the
+    # Anthropic and OpenAI surfaces already expose so V2 callers can reach
+    # the same engine knobs when needed.
+    kv_transfer_params: dict[str, Any] | None = Field(
+        default=None,
+        description="KVTransfer parameters used for disaggregated serving.",
+    )
+    chat_template_kwargs: dict[str, Any] | None = Field(
+        default=None,
+        description=(
+            "Additional keyword args to pass to the chat template renderer. "
+            "Will be accessible by the template."
+        ),
+    )
+
+    @field_validator("model")
+    @classmethod
+    def _validate_model(cls, v: str) -> str:
+        if not v:
+            raise ValueError("model is required")
+        return v
+
+    @field_validator("max_tokens")
+    @classmethod
+    def _validate_max_tokens(cls, v: int | None) -> int | None:
+        if v is not None and v < 0:
+            raise ValueError("max_tokens must be non-negative")
+        return v
+
+    @field_validator("messages", mode="before")
+    @classmethod
+    def _normalize_message_roles(cls, v: Any) -> Any:
+        """Rewrite OpenAI-style ``developer`` roles to ``system``.
+
+        Cohere's v2 ``ChatMessageV2`` discriminated union only admits
+        the four literal roles ``user`` / ``assistant`` / ``system`` /
+        ``tool``. OpenAI's ``developer`` role is documented as
+        high-priority system instructions, so we alias it onto
+        ``system`` *before* the SDK discriminator runs (otherwise it
+        rejects the message with a ``literal_error`` against each union
+        member). Mirrors ``_role_to_melody`` in the renderer so the
+        same alias is honoured no matter which surface the message
+        arrives through.
+
+        ``mode="before"`` is required so the rewrite happens before the
+        ``list[ChatMessageV2]`` coercion runs the SDK's discriminated
+        union; a default-mode validator would never see ``developer``
+        because validation would have already failed. On any
+        structural malformation (non-iterable input, items without a
+        dict-shaped ``role`` field, etc.) we hand ``v`` back unchanged
+        and let Pydantic's normal coercion surface a precise error.
+        """
+        try:
+            return [
+                {**msg, "role": "system"}
+                if msg.get("role", "").lower() == "developer"
+                else msg
+                for msg in v
+            ]
+        except (AttributeError, TypeError):
+            return v
+
+    @field_validator("messages")
+    @classmethod
+    def _validate_messages(cls, v: list[ChatMessageV2]) -> list[ChatMessageV2]:
+        if not v:
+            raise ValueError("messages must contain at least one message")
+        return v
+
+
+# ---------------------------------------------------------------------------
+# Usage / Logprobs
+# ---------------------------------------------------------------------------
+#
+# The Cohere SDK only exposes a v1 ``ApiMetaBilledUnits``. The v2 usage
+# envelope is documented separately in the OpenAPI spec and we declare it
+# here.
+
+
+class CohereUsageBilledUnits(BaseModel):
+    input_tokens: float | None = None
+    output_tokens: float | None = None
+    search_units: float | None = None
+    classifications: float | None = None
+
+
+class CohereUsageTokens(BaseModel):
+    input_tokens: float | None = None
+    output_tokens: float | None = None
+
+
+class CohereUsage(BaseModel):
+    billed_units: CohereUsageBilledUnits | None = None
+    tokens: CohereUsageTokens | None = None
+    cached_tokens: float | None = None
+
+
+class CohereLogprobItem(BaseModel):
+    text: str | None = None
+    token_ids: list[int]
+    logprobs: list[float] | None = None
+
+
+# ---------------------------------------------------------------------------
+# Non-streaming response
+# ---------------------------------------------------------------------------
+
+
+class CohereChatV2Response(BaseModel):
+    """Cohere Chat v2 non-streaming response body.
+
+    Wraps the SDK :class:`AssistantMessageResponse` (the message shape) in
+    the documented v2 response envelope (``id``, ``finish_reason``,
+    ``usage``, ``logprobs``). The single constructor in
+    :class:`CohereServingChatV2._chat_completion_to_v2` is responsible for
+    supplying a non-empty ``id`` (falling back to a synthesized one if
+    the upstream response is missing it) to this model.
+    """
+
+    id: str
+    finish_reason: CohereFinishReason
+    message: AssistantMessageResponse
+    usage: CohereUsage | None = None
+    logprobs: list[CohereLogprobItem] | None = None
+
+    # vLLM-specific extension.
+    kv_transfer_params: dict[str, Any] | None = Field(
+        default=None, description="KVTransfer parameters."
+    )
+
+
+# ---------------------------------------------------------------------------
+# Streaming events
+# ---------------------------------------------------------------------------
+#
+# Cohere V2 streams a sequence of typed JSON events delivered as Server-
+# Sent Events; each event's ``type`` field carries the discriminator. The
+# SDK exposes a Pydantic model per event but none of them declare ``type``
+# as a field (the SDK relies on its own deserializer for discrimination),
+# so a naive ``model_dump_json()`` would silently drop the discriminator
+# and break clients that demux on ``type``.
+#
+# We therefore subclass each SDK event and bake the wire-format ``type``
+# string in as a ``Literal`` field with a default. ``model_dump()`` now
+# emits ``type`` for free, and surfaces can simply construct the event
+# class and serialize it -- no manual ``type`` parameter needed.
+#
+# See https://docs.cohere.com/v2/docs/streaming and the OpenAPI
+# ``StreamedChatResponseV2`` schema for the wire-format reference.
+
+
+class MessageStartEvent(_sdk.ChatMessageStartEvent):
+    type: Literal["message-start"] = "message-start"
+
+
+class ContentStartEvent(_sdk.ChatContentStartEvent):
+    type: Literal["content-start"] = "content-start"
+
+
+class ContentDeltaEvent(_sdk.ChatContentDeltaEvent):
+    type: Literal["content-delta"] = "content-delta"
+
+
+class ContentEndEvent(_sdk.ChatContentEndEvent):
+    type: Literal["content-end"] = "content-end"
+
+
+class ToolPlanDeltaEvent(_sdk.ChatToolPlanDeltaEvent):
+    type: Literal["tool-plan-delta"] = "tool-plan-delta"
+
+
+class ToolCallStartEvent(_sdk.ChatToolCallStartEvent):
+    type: Literal["tool-call-start"] = "tool-call-start"
+
+
+class ToolCallDeltaEvent(_sdk.ChatToolCallDeltaEvent):
+    type: Literal["tool-call-delta"] = "tool-call-delta"
+
+
+class ToolCallEndEvent(_sdk.ChatToolCallEndEvent):
+    type: Literal["tool-call-end"] = "tool-call-end"
+
+
+class CitationStartEvent(_sdk.CitationStartEvent):
+    type: Literal["citation-start"] = "citation-start"
+
+
+class CitationEndEvent(_sdk.CitationEndEvent):
+    type: Literal["citation-end"] = "citation-end"
+
+
+class MessageEndEvent(_sdk.ChatMessageEndEvent):
+    type: Literal["message-end"] = "message-end"

vllm/entrypoints/generate/api_router.py

@@ -4,6 +4,8 @@ from typing import TYPE_CHECKING
 
 from fastapi import FastAPI
 
+import vllm.envs as envs
+
 if TYPE_CHECKING:
     from argparse import Namespace
 
@@ -41,6 +43,12 @@ def register_generate_api_routers(app: FastAPI):
 
     register_anthropic_api_router(app)
 
+    from vllm.entrypoints.cohere.api_router import (
+        attach_router as register_cohere_api_router,
+    )
+
+    register_cohere_api_router(app)
+
     from .generative_scoring.api_router import register_generative_scoring_api_router
 
     register_generative_scoring_api_router(app)
@@ -55,6 +63,23 @@ async def init_generate_state(
 ):
     from vllm.entrypoints.anthropic.serving import AnthropicServingMessages
     from vllm.entrypoints.chat_utils import load_chat_template
+
+    # The Cohere serving handler depends on the optional `cohere` SDK for
+    # its wire-format protocol models, and is additionally gated on the
+    # `VLLM_ENABLE_COHERE_API` env flag (see
+    # `vllm.entrypoints.cohere.api_router.attach_router`). Skip the import
+    # entirely when the endpoint isn't going to be exposed, both because
+    # the SDK may not be installed and because the serving object holds
+    # nontrivial state (chat handler, warmup) that would otherwise be
+    # unused.
+    if envs.VLLM_ENABLE_COHERE_API:
+        try:
+            from vllm.entrypoints.cohere.serving import CohereServingChatV2
+        except ImportError:
+            CohereServingChatV2 = None  # type: ignore[assignment,misc]
+    else:
+        CohereServingChatV2 = None  # type: ignore[assignment,misc]
+
     from vllm.entrypoints.mcp.tool_server import (
         DemoToolServer,
         MCPToolServer,
@@ -86,6 +111,19 @@ async def init_generate_state(
         tool_server = None
     resolved_chat_template = load_chat_template(args.chat_template)
 
+    # Fold the dedicated ``--cohere-format`` CLI flag into the renderer's
+    # default chat-template kwargs. The cohere renderer reads
+    # ``chat_template_kwargs["cohere_format"]`` to pick cmd3 vs cmd4
+    # rendering; making this a first-class flag keeps the right format
+    # discoverable for ``vllm serve --tokenizer-mode cohere`` users
+    # without forcing them to hand-construct a JSON dict for
+    # ``--default-chat-template-kwargs``. Per-request overrides still
+    # take precedence (see ``merge_kwargs`` in
+    # ``ChatCompletionRequest.build_chat_params``).
+    default_chat_template_kwargs = dict(args.default_chat_template_kwargs or {})
+    if getattr(args, "cohere_format", None):
+        default_chat_template_kwargs.setdefault("cohere_format", args.cohere_format)
+
     # Render endpoints are always backed by OnlineRenderer so that
     # /v1/chat/completions/render and /v1/completions/render work on both
     # generate-mode and render-only servers. Created in init_app_state.
@@ -106,7 +144,7 @@ async def init_generate_state(
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
             enable_log_outputs=args.enable_log_outputs,
-            default_chat_template_kwargs=args.default_chat_template_kwargs,
+            default_chat_template_kwargs=default_chat_template_kwargs,
         )
         if "generate" in supported_tasks
         else None
@@ -119,7 +157,7 @@ async def init_generate_state(
         request_logger=request_logger,
         chat_template=resolved_chat_template,
         chat_template_content_format=args.chat_template_content_format,
-        default_chat_template_kwargs=args.default_chat_template_kwargs,
+        default_chat_template_kwargs=default_chat_template_kwargs,
         trust_request_chat_template=args.trust_request_chat_template,
         return_tokens_as_token_ids=args.return_tokens_as_token_ids,
         enable_auto_tools=args.enable_auto_tool_choice,
@@ -169,11 +207,32 @@ async def init_generate_state(
             reasoning_parser=args.structured_outputs_config.reasoning_parser,
             enable_prompt_tokens_details=args.enable_prompt_tokens_details,
             enable_force_include_usage=args.enable_force_include_usage,
-            default_chat_template_kwargs=args.default_chat_template_kwargs,
+            default_chat_template_kwargs=default_chat_template_kwargs,
         )
         if "generate" in supported_tasks
         else None
     )
+    state.cohere_serving_chat_v2 = (
+        CohereServingChatV2(
+            engine_client,
+            state.openai_serving_models,
+            args.response_role,
+            online_renderer=state.online_renderer,
+            request_logger=request_logger,
+            chat_template=resolved_chat_template,
+            chat_template_content_format=args.chat_template_content_format,
+            return_tokens_as_token_ids=args.return_tokens_as_token_ids,
+            enable_auto_tools=args.enable_auto_tool_choice,
+            tool_parser=args.tool_call_parser,
+            reasoning_parser=args.structured_outputs_config.reasoning_parser,
+            enable_prompt_tokens_details=args.enable_prompt_tokens_details,
+            enable_force_include_usage=args.enable_force_include_usage,
+            default_chat_template_kwargs=default_chat_template_kwargs,
+            is_reasoning_model=args.cohere_is_reasoning_model,
+        )
+        if CohereServingChatV2 is not None and "generate" in supported_tasks
+        else None
+    )
 
     from .generative_scoring.serving import ServingGenerativeScoring
 

vllm/entrypoints/openai/chat_completion/protocol.py

@@ -10,7 +10,13 @@ from openai.types.chat.chat_completion_audio import (
     ChatCompletionAudio as OpenAIChatCompletionAudio,
 )
 from openai.types.chat.chat_completion_message import Annotation as OpenAIAnnotation
-from pydantic import Field, PrivateAttr, model_serializer, model_validator
+from pydantic import (
+    Field,
+    PrivateAttr,
+    SerializeAsAny,
+    model_serializer,
+    model_validator,
+)
 
 from vllm.config import ModelConfig
 from vllm.entrypoints.chat_utils import (
@@ -91,7 +97,12 @@ class ChatCompletionLogProbs(OpenAIBaseModel):
 
 class ChatCompletionResponseChoice(OpenAIBaseModel):
     index: int
-    message: ChatMessage
+    # ``SerializeAsAny`` lets pydantic honor subclasses of ``ChatMessage``
+    # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereChatMessage``)
+    # so that added fields like ``citations`` survive JSON serialization
+    # instead of being stripped down to the base schema. Plain
+    # ``ChatMessage`` instances serialize identically to before.
+    message: SerializeAsAny[ChatMessage]
     logprobs: ChatCompletionLogProbs | None = None
     # per OpenAI spec this is the default
     finish_reason: str | None = "stop"
@@ -139,7 +150,11 @@ class ChatCompletionResponse(OpenAIBaseModel):
 
 class ChatCompletionResponseStreamChoice(OpenAIBaseModel):
     index: int
-    delta: DeltaMessage
+    # ``SerializeAsAny`` lets pydantic honor subclasses of ``DeltaMessage``
+    # (e.g. ``vllm.entrypoints.cohere.cohere_chat_message.CohereDeltaMessage``)
+    # so streaming ``citations`` survive JSON serialization. Plain
+    # ``DeltaMessage`` instances serialize identically to before.
+    delta: SerializeAsAny[DeltaMessage]
     logprobs: ChatCompletionLogProbs | None = None
     finish_reason: str | None = None
     stop_reason: int | str | None = None

vllm/entrypoints/openai/chat_completion/serving.py

@@ -391,6 +391,34 @@ class OpenAIServingChat(GenerateBaseServing):
             return self.response_role
         return request.messages[-1]["role"]
 
+    def _create_chat_message(self, *args: Any, **kwargs: Any) -> ChatMessage:
+        """Construct the response :class:`ChatMessage` for the non-streaming path.
+
+        The full-generator calls this at every construction site so
+        subclasses can swap in a specialized :class:`ChatMessage`
+        subclass (e.g. :class:`CohereServingChatV2` returning
+        :class:`CohereChatMessage`) without duplicating the branchy
+        tool-choice / auto-tools logic that decides which fields are
+        populated. The default returns a plain :class:`ChatMessage`.
+        """
+        return ChatMessage(*args, **kwargs)
+
+    def _finalize_response_message(
+        self,
+        message: ChatMessage,
+        *,
+        parser: Parser | None,
+    ) -> ChatMessage:
+        """Subclass hook to enrich a fully-constructed :class:`ChatMessage`.
+
+        Default is a no-op. Subclasses that need to surface parser-side
+        extras (e.g. :class:`CohereServingChatV2` reading grounding
+        citations off the reasoning parser and populating
+        :class:`CohereChatMessage.citations`) override this to inspect
+        ``parser`` and mutate/replace ``message``.
+        """
+        return message
+
     async def chat_completion_stream_generator(
         self,
         request: ChatCompletionRequest,
@@ -894,13 +922,19 @@ class OpenAIServingChat(GenerateBaseServing):
             )
             is_required_tool_choice = request.tool_choice == "required"
 
+            # All six construction sites route through ``self._create_chat_message``
+            # so subclasses can swap in a specialized :class:`ChatMessage`
+            # (e.g. the Cohere v2 handler's ``CohereChatMessage``) without
+            # having to duplicate this branch logic.
             if (not self.enable_auto_tools or not tool_parser_cls) and (
                 not is_named_tool_choice and not is_required_tool_choice
             ):
-                message = ChatMessage(role=role, reasoning=reasoning, content=content)
+                message = self._create_chat_message(
+                    role=role, reasoning=reasoning, content=content
+                )
 
             elif is_named_tool_choice or is_required_tool_choice:
-                message = ChatMessage(
+                message = self._create_chat_message(
                     role=role,
                     reasoning=reasoning,
                     content=content or "",
@@ -913,7 +947,9 @@ class OpenAIServingChat(GenerateBaseServing):
             # if the request doesn't use tool choice
             # OR specifies to not use a tool
             elif not request.tool_choice or request.tool_choice == "none":
-                message = ChatMessage(role=role, reasoning=reasoning, content=content)
+                message = self._create_chat_message(
+                    role=role, reasoning=reasoning, content=content
+                )
 
             # handle when there are tools and tool choice is auto
             elif (
@@ -924,7 +960,7 @@ class OpenAIServingChat(GenerateBaseServing):
             ):
                 auto_tools_called = tool_calls is not None and len(tool_calls) > 0
                 if tool_calls:
-                    message = ChatMessage(
+                    message = self._create_chat_message(
                         role=role,
                         reasoning=reasoning,
                         content=content,
@@ -935,7 +971,7 @@ class OpenAIServingChat(GenerateBaseServing):
                     )
 
                 else:
-                    message = ChatMessage(
+                    message = self._create_chat_message(
                         role=role,
                         reasoning=reasoning,
                         content=content,
@@ -948,7 +984,17 @@ class OpenAIServingChat(GenerateBaseServing):
                     " if tools should be extracted. Returning a standard chat "
                     "completion."
                 )
-                message = ChatMessage(role=role, reasoning=reasoning, content=content)
+                message = self._create_chat_message(
+                    role=role, reasoning=reasoning, content=content
+                )
+
+            # Subclass hook: enrich the constructed message with any
+            # parser-side extras that don't fit through the plain
+            # ``(reasoning, content, tool_calls)`` tuple. Base is a no-op;
+            # citation-aware handlers use this to surface grounding
+            # metadata cached on the reasoning parser.
+            message = self._finalize_response_message(message, parser=parser)
+
             # In OpenAI's API, when a tool is called, the finish_reason is:
             # "tool_calls" for "auto" or "required" tool calls,
             # and "stop" for named tool calls.

vllm/entrypoints/openai/cli_args.py

@@ -148,6 +148,25 @@ class BaseFrontendArgs:
     """If set to False, output deltas will not be logged. Relevant only if 
     --enable-log-outputs is set.
     """
+    cohere_is_reasoning_model: bool = True
+    """Cohere ``/cohere/v2/chat`` only. Whether the served model is a
+    reasoning Command-family model. When True (default), the assistant's
+    chain-of-thought is surfaced as a ``thinking`` content block (or
+    ``content-*`` events on the stream). When False, reasoning is
+    surfaced as Cohere's ``tool_plan`` field (or ``tool-plan-delta``
+    events) whenever the model emits tool calls, matching older non-
+    reasoning Command models. Has no effect on the non-Cohere
+    endpoints."""
+    cohere_format: str = "cmd4"
+    """Cohere ``--tokenizer-mode cohere`` only. Which Cohere prompt
+    format to render: ``cmd4`` (current Command A+ models; default) or
+    ``cmd3`` (earlier Cmd-A and Cmd-A reasoning models). Selecting the
+    wrong format silently produces a prompt the model wasn't trained
+    on, which most commonly manifests as the model emitting text but no
+    citations / tool calls / thinking blocks. Equivalent to passing
+    ``--default-chat-template-kwargs '{"cohere_format": "..."}'`` -- any
+    explicit request-level ``chat_template_kwargs.cohere_format`` takes
+    priority."""
     log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE
     """If set to True, log the stack trace of error responses"""
     tokens_only: bool = False

vllm/entrypoints/serve/utils/server_utils.py

@@ -39,7 +39,7 @@ from vllm.v1.engine.exceptions import EngineDeadError, EngineGenerateError
 logger = init_logger("vllm.entrypoints.openai.server_utils")
 
 
-GUARDED_PREFIX = ("/v1", "/v2", "/inference")
+GUARDED_PREFIX = ("/v1", "/v2", "/inference", "/cohere")
 
 
 class AuthenticationMiddleware:

vllm/envs.py

@@ -240,6 +240,7 @@ if TYPE_CHECKING:
     VLLM_LOOPBACK_IP: str = ""
     VLLM_ALLOW_CHUNKED_LOCAL_ATTN_WITH_HYBRID_KV_CACHE: bool = True
     VLLM_ENABLE_RESPONSES_API_STORE: bool = False
+    VLLM_ENABLE_COHERE_API: bool = False
     VLLM_HAS_FLASHINFER_CUBIN: bool = False
     VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
@@ -1758,6 +1759,11 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ENABLE_RESPONSES_API_STORE": lambda: bool(
         int(os.getenv("VLLM_ENABLE_RESPONSES_API_STORE", "0"))
     ),
+    # If set to 1, expose the Cohere Chat v2 API at ``POST /cohere/v2/chat``.
+    # Default off
+    "VLLM_ENABLE_COHERE_API": lambda: bool(
+        int(os.getenv("VLLM_ENABLE_COHERE_API", "0"))
+    ),
     # If set, use the fp8 mfma in rocm paged attention.
     "VLLM_ROCM_FP8_MFMA_PAGE_ATTN": lambda: bool(
         int(os.getenv("VLLM_ROCM_FP8_MFMA_PAGE_ATTN", "0"))

vllm/reasoning/cohere_command_reasoning_parser.py

@@ -20,6 +20,11 @@ except ImportError as e:
     ) from e
 
 
+from vllm.entrypoints.cohere.cohere_chat_message import (
+    Citation,
+    CitationSource,
+    CohereDeltaMessage,
+)
 from vllm.entrypoints.openai.chat_completion.protocol import (
     ChatCompletionRequest,
 )
@@ -404,6 +409,83 @@ def _schema_dict_from_structured_outputs(
     )
 
 
+def _melody_sources_to_vllm(
+    raw_sources: Any,
+    position_to_source: Mapping[tuple[int, int], CitationSource] | None,
+) -> list[CitationSource]:
+    """Convert melody's ``Source`` objects into resolved :class:`CitationSource`.
+
+    Melody's ``Source`` shape is
+    ``{tool_call_index, tool_result_indices, document_ids}``. Without
+    :meth:`PyFilterOptions.with_message_history` the ``document_ids``
+    list is empty, so we resolve the numeric address ourselves using
+    ``position_to_source``, which is
+    :meth:`vllm.entrypoints.cohere.serving.CohereServingChatV2._build_position_to_source`
+    applied to the inbound request and forwarded through
+    ``chat_template_kwargs`` at parser-construction time. Each melody
+    source with ``tool_result_indices=[i, j, ...]`` fans out to one
+    :class:`CitationSource` per resolved position, populating the
+    wire-shape ``type`` / ``id`` / ``document`` / ``tool_output``
+    fields; unresolved positions are silently skipped (the calling
+    citation is dropped at wire-coercion time if every source in it
+    ends up unresolved -- see ``_to_wire_citation`` in
+    :mod:`vllm.entrypoints.cohere.serving`).
+
+    When ``position_to_source`` is ``None`` (parser wired outside of
+    :class:`CohereServingChatV2`, e.g. a plain OpenAI chat completion
+    request that happens to hit a cohere reasoning parser) sources
+    can't be resolved and get dropped; citations then never reach any
+    citation-aware handler downstream.
+    """
+    if not position_to_source:
+        return []
+    out: list[CitationSource] = []
+    for s in raw_sources or []:
+        bucket = getattr(s, "tool_call_index", None)
+        indices = list(getattr(s, "tool_result_indices", None) or [])
+        if bucket is None or not indices:
+            continue
+        for idx in indices:
+            info = position_to_source.get((bucket, idx))
+            if info is not None:
+                out.append(info)
+    return out
+
+
+def _melody_citations_to_vllm(
+    raw_citations: Any,
+    position_to_source: Mapping[tuple[int, int], CitationSource] | None,
+) -> list[Citation] | None:
+    """Convert melody's ``FilterCitation`` objects into :class:`Citation`.
+
+    Resolves each source in-place via ``position_to_source`` (see
+    :func:`_melody_sources_to_vllm`). Citations whose sources all fail
+    to resolve are still emitted here with ``sources=[]``; the serving
+    layer drops them at wire-coercion time so the fail-closed policy
+    lives in exactly one place.
+    """
+    if not raw_citations:
+        return None
+    out: list[Citation] = []
+    for c in raw_citations:
+        out.append(
+            Citation(
+                start=getattr(c, "start_index", None),
+                end=getattr(c, "end_index", None),
+                text=getattr(c, "text", None),
+                sources=_melody_sources_to_vllm(
+                    getattr(c, "sources", None), position_to_source
+                ),
+                type=(
+                    "THINKING_CONTENT"
+                    if getattr(c, "is_thinking", False)
+                    else "TEXT_CONTENT"
+                ),
+            )
+        )
+    return out
+
+
 class BaseCohereCommandReasoningParser(ReasoningParser):
     def __init__(
         self,
@@ -420,6 +502,32 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
         self.unary_opts = unary_opts
         self.melody_unary = PyFilter(unary_opts)
         self.melody_streaming = PyFilter(streaming_opts)
+        # Citations extracted by the most recent ``extract_reasoning`` call.
+        # Citation-aware handlers (e.g. :class:`CohereServingChatV2` via its
+        # ``_finalize_response_message`` hook) read this back from the
+        # parser instance (which is constructed per-request) and attach
+        # the result to :class:`CohereChatMessage.citations`. ``None`` when
+        # the last parse produced no citations.
+        self.last_unary_citations: list[Citation] | None = None
+        # Request-scoped ``(tool_call_index, tool_result_idx) -> resolved
+        # CitationSource`` map, forwarded from
+        # :meth:`CohereServingChatV2._apply_cohere_template_kwargs` via
+        # ``chat_template_kwargs``. See :func:`_melody_sources_to_vllm`
+        # for how it's consumed. Absent for parser instances that
+        # weren't created by ``CohereServingChatV2`` (e.g. an OpenAI
+        # chat completion request routed through a cohere reasoning
+        # parser); citations from those requests are dropped since we
+        # have no way to attribute their sources.
+        ctk = kwargs.get("chat_template_kwargs") or {}
+        # Lazy import: ``renderers.cohere`` transitively imports
+        # ``cohere_melody`` and other heavy deps; hoisting the constant
+        # to module scope would tie parser importability to that graph.
+        from vllm.renderers.cohere import POSITION_TO_SOURCE_KEY
+
+        raw_map = ctk.get(POSITION_TO_SOURCE_KEY)
+        self._position_to_source: Mapping[tuple[int, int], CitationSource] | None = (
+            raw_map if isinstance(raw_map, Mapping) else None
+        )
 
     @property
     def reasoning_start_str(self) -> str | None:
@@ -439,9 +547,22 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
         delta_token_ids: Sequence[int],
     ) -> DeltaMessage | None:
         r = self.melody_streaming.write_decoded(delta_text)
-        if r.content is None and r.reasoning is None and not r.tool_calls:
+        citations = _melody_citations_to_vllm(
+            getattr(r, "citations", None), self._position_to_source
+        )
+        if (
+            r.content is None
+            and r.reasoning is None
+            and not r.tool_calls
+            and not citations
+        ):
             return None
-        msg = DeltaMessage()
+        # Always emit CohereDeltaMessage so citations reach the wire via
+        # ``SerializeAsAny[DeltaMessage]`` on the streaming choice envelope.
+        # When ``citations`` is unset, ``CohereDeltaMessage._serialize``
+        # drops the field so the emitted shape matches plain
+        # :class:`DeltaMessage`
+        msg = CohereDeltaMessage()
         if r.content is not None:
             msg.content = r.content
         if r.reasoning is not None:
@@ -456,12 +577,23 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
                 )
                 for tc in r.tool_calls
             ]
+        if citations:
+            msg.citations = citations
         return msg
 
     def extract_reasoning(
         self, model_output: str, request: ChatCompletionRequest | ResponsesRequest
     ) -> tuple[str | None, str | None]:
         result = self.melody_unary.process_full_text(model_output)
+        # Cache citations so citation-aware handlers can surface them
+        # on :class:`CohereChatMessage.citations`. The base
+        # :meth:`ReasoningParser.extract_reasoning` contract only
+        # returns ``(reasoning, content)``, so citations are passed
+        # back via parser-instance state -- safe because the parser is
+        # constructed per-request.
+        self.last_unary_citations = _melody_citations_to_vllm(
+            getattr(result, "citations", None), self._position_to_source
+        )
         return result.reasoning, result.content
 
     def extract_content_ids(self, input_ids: list[int]) -> list[int]:
@@ -549,12 +681,25 @@ class BaseCohereCommandReasoningParser(ReasoningParser):
         return request
 
 
+# melody's streaming filter only buffers a partial ``<co: ...>`` citation
+# across ``write_decoded`` calls when ``stream_non_grounded_answer`` is
+# set: otherwise, the moment an opening ``<co`` is seen without a
+# closing ``</co: ...>`` in the same delta, the filter emits the partial
+# marker bytes verbatim as plain content. In vLLM's streaming path the
+# parser is fed one token (1-4 chars) per call, so an unbuffered filter
+# will leak ``<co: 0>``-style markers into ``delta.content`` and never
+# emit a ``FilterCitation`` for them. Enabling the flag flips the
+# partial-match branch in melody's ``parse_citations`` (see
+# ``src/parsing/citations_filter.rs``) to ``return (None, 0)`` -- i.e.
+# keep buffering -- which lets a full citation eventually resolve.
+# Non-streaming (unary) parsing receives the whole output in one call so
+# the flag is a no-op there and we leave ``unary_opts`` alone.
 class CohereCommand3ReasoningParser(BaseCohereCommandReasoningParser):
     def __init__(self, tokenizer: TokenizerLike, *args, **kwargs):
         super().__init__(
             tokenizer,
             *args,
-            streaming_opts=PyFilterOptions().cmd3(),
+            streaming_opts=PyFilterOptions().cmd3().stream_non_grounded_answer(),
             unary_opts=PyFilterOptions().cmd3().no_tools(),
             **kwargs,
         )
@@ -565,7 +710,7 @@ class CohereCommand4ReasoningParser(BaseCohereCommandReasoningParser):
         super().__init__(
             tokenizer,
             *args,
-            streaming_opts=PyFilterOptions().cmd4(),
+            streaming_opts=PyFilterOptions().cmd4().stream_non_grounded_answer(),
             unary_opts=PyFilterOptions().cmd4().no_tools(),
             **kwargs,
         )

vllm/renderers/cohere.py

@@ -0,0 +1,704 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Cohere prompt renderer.
+
+Templates the Cohere Command-family prompt formats (cmd3 / cmd4) using
+the ``cohere_melody`` Rust bindings instead of Jinja. Enabled by
+passing ``--tokenizer-mode cohere`` on the engine; tokenization itself
+still flows through the cached HuggingFace tokenizer.
+
+The renderer's client surface is ``chat_template_kwargs`` (vLLM's
+generic passthrough for template-time inputs). It is fed by two
+different code paths, which speak different vocabularies:
+
+1. :class:`vllm.entrypoints.cohere.serving.CohereServingChatV2` --
+   translates a native `Cohere Chat v2 request
+   <https://docs.cohere.com/reference/chat>`__ into
+   ``chat_template_kwargs``.
+2. :class:`OpenAIServingChat` -- forwards whatever the client puts in
+   ``request.chat_template_kwargs`` directly, without translation.
+
+Accordingly the renderer accepts keys from both vocabularies.
+
+**Cohere Chat v2 request fields** (populated automatically by
+``CohereServingChatV2``; a direct ``chat_template_kwargs`` caller may
+also pass them):
+
+* ``documents``: list of documents (v2 ``documents``)
+* ``tools``: list of tools (v2 ``tools``)
+* ``safety_mode``: v2 ``safety_mode`` (``CONTEXTUAL`` / ``STRICT`` /
+  ``NONE``) -- forwarded to melody's cmd3 ``safety_mode`` slot;
+  lowercased.
+* ``citation_options``: v2 ``citation_options``. Its ``.mode``
+  (``ENABLED`` / ``DISABLED`` / ``FAST`` / ``ACCURATE`` / ``OFF``) is
+  normalized into melody's cmd3 ``citation_quality`` (``on`` / ``off``)
+  and cmd4 ``grounding`` (``enabled`` / ``disabled`` / ``unknown``)
+  slots. cmd4 has no fast/accurate distinction at the prompt-template
+  layer -- both mean grounding-on.
+* ``response_format``: v2 ``response_format`` (``json_object`` /
+  ``json_schema``) -- mapped onto melody's ``json_mode`` /
+  ``json_schema`` config slots.
+* ``thinking``: v2 ``thinking``; its ``.type`` (``enabled`` /
+  ``disabled``) becomes melody's ``reasoning_type``.
+
+**Melody template-config knobs** (accepted for direct-passthrough
+callers; not part of the Cohere Chat v2 surface):
+
+* ``cohere_format``: renderer selector, ``cmd3`` or ``cmd4`` (default
+  ``cmd4``).
+* ``template_id``: pick one of melody's built-in template variants.
+* ``available_tools``: melody's own name for tools; takes precedence
+  over ``tools`` when both are set.
+* ``reasoning_type``: direct ``enabled`` / ``disabled`` toggle;
+  overrides derivation from ``thinking.type``.
+* ``dev_instruction``: developer-instruction override.
+* ``json_mode`` / ``json_schema``: direct structured-output toggles;
+  override derivation from ``response_format``.
+* cmd3-only: ``citation_quality`` (``on`` / ``off``), ``skip_preamble``.
+* cmd4-only: ``grounding`` (``enabled`` / ``disabled`` / ``unknown``;
+  overrides derivation from ``citation_options.mode``),
+  ``platform_instruction``.
+
+**Rejected inputs**: ``template_jinja`` and ``template`` are Cohere's
+own inlets for raw Jinja template source. They are valid at Cohere's
+API surface but not at vLLM's -- raw template source must flow through
+the standard ``chat_template`` request field (guarded by
+``--trust-request-chat-template``), which this renderer forwards to
+melody as ``template_jinja``. Setting these keys explicitly in
+``chat_template_kwargs`` raises ``ValueError`` so client
+misconfiguration surfaces loudly.
+
+**Everything else** in ``chat_template_kwargs`` is forwarded verbatim
+to melody as ``additional_template_fields`` and becomes accessible as
+Jinja variables inside the template. This matches vLLM's documented
+contract for ``chat_template_kwargs`` ("kwargs accessible by the
+template"), so e.g. ``chat_template_kwargs={"reasoning_effort":
+"low"}`` resolves ``{{ reasoning_effort }}`` inside cmd3 / cmd4
+templates.
+
+Citations produced by Cohere models are surfaced through the
+Cohere-scoped ``CohereChatMessage.citations`` /
+``CohereDeltaMessage.citations`` fields (see
+:mod:`vllm.entrypoints.cohere.cohere_chat_message`), populated by the
+``cohere2`` reasoning parser. The base OpenAI ``ChatMessage`` /
+``DeltaMessage`` keep their declared schemas unchanged; the response
+envelope declares them as ``SerializeAsAny[...]`` so the subclass
+fields survive JSON serialization.
+"""
+
+from __future__ import annotations
+
+import copy
+import json
+from enum import Enum
+from typing import Any
+
+from vllm.config import VllmConfig
+from vllm.entrypoints.chat_utils import (
+    ChatCompletionMessageParam,
+    ConversationMessage,
+    parse_chat_messages,
+    parse_chat_messages_async,
+)
+from vllm.logger import init_logger
+from vllm.tokenizers.hf import HfTokenizer
+from vllm.utils.async_utils import make_async
+
+from .base import BaseRenderer
+from .inputs import DictPrompt
+from .inputs.preprocess import parse_dec_only_prompt
+from .params import ChatParams
+
+logger = init_logger(__name__)
+
+
+_DEFAULT_FORMAT = "cmd4"
+_VALID_FORMATS = ("cmd3", "cmd4")
+
+
+class MelodyContentType(str, Enum):
+    """Wire-format discriminator for melody content blocks.
+
+    These strings are what the cmd3 / cmd4 Jinja templates check against
+    in ``message.content[0].type`` (e.g. the ``thinking`` branch in
+    ``cmd4-v1.jinja``). Keep new values in sync with melody's template
+    schema.
+    """
+
+    TEXT = "text"
+    THINKING = "thinking"
+    IMAGE = "image"
+    DOCUMENT = "document"
+
+
+# Reserved ``chat_template_kwargs`` key that carries
+# ``AssistantChatMessageV2.citations`` from a Cohere v2 request through
+# to the renderer, keyed by request-message index. Populated by
+# ``CohereServingChatV2._apply_cohere_template_kwargs`` and consumed by
+# ``_conversation_to_melody_messages``. The value is a
+# ``dict[int, list[dict]]`` of melody-shape ``FilterCitation`` dicts.
+# The leading underscore signals "internal, not part of the public
+# ``chat_template_kwargs`` surface" -- clients should not set this key
+# directly.
+MESSAGES_CITATIONS_KEY = "_messages_citations"
+
+
+# Reserved ``chat_template_kwargs`` key carrying a request-side
+# ``dict[tuple[int, int], CitationSource]`` map from melody's numeric
+# ``(tool_call_index, tool_result_index)`` addressing to fully-resolved
+# wire-shape sources (type, id, document/tool_output payload).
+# Populated by :meth:`CohereServingChatV2._apply_cohere_template_kwargs`
+# and consumed by :class:`BaseCohereCommandReasoningParser` at request
+# start so the parser can attach real document ids on each melody source
+# it emits -- eliminating the numeric-coord round-trip through the
+# internal streaming pipeline. Leading underscore signals "internal,
+# not part of the public ``chat_template_kwargs`` surface" -- clients
+# should not set this key directly.
+POSITION_TO_SOURCE_KEY = "_position_to_source"
+
+
+# Keys this renderer interprets directly from ``chat_template_kwargs`` and
+# maps onto typed melody render-config fields. Everything *not* in this
+# set is forwarded verbatim to melody as ``additional_template_fields``
+# (i.e. as Jinja template variables), so callers can write
+# ``chat_template_kwargs = {"my_var": "..."}`` and have ``{{ my_var }}``
+# resolve inside the template -- matching vLLM's documented contract
+# for ``chat_template_kwargs``.
+_RENDERER_CONSUMED_KEYS = frozenset(
+    {
+        "cohere_format",
+        "template_id",
+        # ``template_jinja`` / ``template`` are Cohere-only inputs for
+        # raw Jinja source -- valid at Cohere's own API surface but not
+        # at vLLM's. In vLLM, raw template source must flow through the
+        # standard ``chat_template`` request field, so a non-``None``
+        # value here is rejected in ``_build_render_config``. The keys
+        # stay in the consumed set so an explicit ``None`` (e.g. from
+        # ``.model_dump(exclude_none=False)``) is dropped rather than
+        # surfacing as a stray ``{{ template_jinja }}`` Jinja variable.
+        "template_jinja",
+        "template",
+        "use_jinja",
+        "documents",
+        "available_tools",
+        "tools",
+        "reasoning_type",
+        "thinking",
+        "dev_instruction",
+        "response_format",
+        "json_schema",
+        "json_mode",
+        "safety_mode",
+        "citation_quality",
+        "citation_options",
+        "skip_preamble",
+        "grounding",
+        "platform_instruction",
+        MESSAGES_CITATIONS_KEY,
+        POSITION_TO_SOURCE_KEY,
+    }
+)
+
+
+def _try_import_melody():
+    try:
+        import cohere_melody  # type: ignore
+
+        return cohere_melody
+    except ImportError as e:  # pragma: no cover - exercised at runtime
+        raise ImportError(
+            "The `cohere` tokenizer/renderer mode requires the "
+            "`cohere_melody` package. Install it via "
+            "`pip install cohere-melody` or build from "
+            "https://github.com/cohere-ai/melody."
+        ) from e
+
+
+_MELODY_ROLES = frozenset({"system", "user", "chatbot", "tool"})
+
+
+# Cohere v2 ``citation_options.mode`` -> melody cmd4 ``grounding``.
+# v2 surfaces five modes (``ENABLED`` / ``DISABLED`` / ``FAST`` /
+# ``ACCURATE`` / ``OFF``) but cmd4's template doesn't differentiate
+# fast vs accurate at the prompt layer -- both just turn grounding on,
+# and ``ENABLED`` / ``DISABLED`` are direct aliases. Melody itself
+# accepts ``unknown`` / ``enabled`` / ``disabled``.
+_CMD4_GROUNDING_FROM_MODE = {
+    "fast": "enabled",
+    "accurate": "enabled",
+    "on": "enabled",
+    "enabled": "enabled",
+    "off": "disabled",
+    "disabled": "disabled",
+    "unknown": "unknown",
+}
+
+
+def _normalize_cmd4_grounding(value: Any) -> str:
+    """Coerce a user-facing grounding/citation mode to melody's vocab.
+
+    Raises ``ValueError`` rather than letting an unrecognized value
+    slip through to ``render_cmd4`` and surface as a generic
+    ``Invalid config: grounding`` from melody.
+    """
+    out = _CMD4_GROUNDING_FROM_MODE.get(value.lower())
+    if out is None:
+        raise ValueError(
+            f"Unrecognized cmd4 grounding value: {value!r}. Expected one "
+            f"of ENABLED / DISABLED / FAST / ACCURATE / OFF (from "
+            f"citation_options.mode), or the direct melody values "
+            f"enabled / disabled / unknown."
+        )
+    return out
+
+
+def _role_to_melody(role: str) -> str:
+    """Map an OpenAI role to the role string melody expects.
+
+    The cmd3 / cmd4 jinja templates only recognize ``system``, ``user``,
+    ``assistant``/``chatbot``, and ``tool``: any other role is silently
+    dropped by the template's role-dispatch chain (no fallback branch),
+    which produces a malformed prompt without any error. We therefore
+    refuse unknown roles up front rather than letting them disappear.
+
+    Aliases:
+
+    * ``assistant`` -> ``chatbot`` (Cohere's historical assistant role
+      name used by the templates).
+    * ``developer`` -> ``system`` (OpenAI's ``developer`` role is
+      documented as high-priority instructions, which maps onto the
+      ``system`` slot in Cohere's prompt format).
+    """
+    role = role.lower()
+    if role == "assistant":
+        return "chatbot"
+    if role == "developer":
+        return "system"
+    if role in _MELODY_ROLES:
+        return role
+    raise ValueError(
+        f"Unsupported message role for the cohere renderer: {role!r}. "
+        "Expected one of: system, developer, user, assistant, chatbot, tool."
+    )
+
+
+def _normalize_tool_call(tc: dict[str, Any] | Any) -> dict[str, Any]:
+    """Normalize an OpenAI tool call dict into melody's tool call shape.
+
+    melody expects ``{id, name, parameters: <json str>}`` whereas OpenAI
+    delivers ``{id, type, function: {name, arguments: <json|str>}}``.
+    """
+    # Pydantic objects
+    if hasattr(tc, "model_dump"):
+        tc = tc.model_dump()
+    if not isinstance(tc, dict):
+        raise TypeError(f"Unexpected tool_call value: {tc!r}")
+
+    fn = tc.get("function") or {}
+    name = fn.get("name") or tc.get("name") or ""
+    args = fn.get("arguments")
+    if args is None:
+        args = tc.get("arguments", {})
+    # melody expects a JSON-encoded string
+    if not isinstance(args, str):
+        args = json.dumps(args, ensure_ascii=False)
+    return {
+        "id": tc.get("id") or "",
+        "name": name,
+        "parameters": args,
+    }
+
+
+def _content_blocks(content: Any) -> list[dict[str, Any]]:
+    """Convert OpenAI ``ConversationMessage.content`` into melody content blocks.
+
+    The chat_utils ``content_format="openai"`` produces a list of
+    ``{type: text, text: ...}`` dicts already; we pass those through with
+    minimal coercion. Plain string content is wrapped in a single text block.
+    Image / multimodal placeholder dicts (``{"type": "image"}``) are
+    forwarded as-is so that templates can reference image placeholders that
+    the upstream tokenizer expands separately.
+    """
+    if content is None:
+        return []
+    if isinstance(content, str):
+        return [{"type": MelodyContentType.TEXT, "text": content}]
+    blocks: list[dict[str, Any]] = []
+    for part in content:
+        if isinstance(part, str):
+            blocks.append({"type": MelodyContentType.TEXT, "text": part})
+            continue
+        if not isinstance(part, dict):
+            raise TypeError(f"Unexpected content part: {part!r}")
+        part_type = part.get("type", MelodyContentType.TEXT)
+        if part_type in ("text", "input_text", "output_text", "refusal"):
+            blocks.append(
+                {"type": MelodyContentType.TEXT, "text": part.get("text", "")}
+            )
+        elif part_type == MelodyContentType.THINKING:
+            blocks.append(
+                {
+                    "type": MelodyContentType.THINKING,
+                    "thinking": part.get("thinking", ""),
+                }
+            )
+        elif part_type == MelodyContentType.IMAGE:
+            blocks.append(
+                {
+                    "type": MelodyContentType.IMAGE,
+                    "image": {
+                        "template_placeholder": part.get(
+                            "template_placeholder", "<image>"
+                        ),
+                    },
+                }
+            )
+        elif part_type == MelodyContentType.DOCUMENT:
+            doc = part.get("document")
+            if isinstance(doc, dict):
+                blocks.append({"type": MelodyContentType.DOCUMENT, "document": doc})
+            else:
+                # Fall back to wrapping arbitrary string as text.
+                blocks.append({"type": MelodyContentType.TEXT, "text": json.dumps(doc)})
+        elif part_type == "tool_reference":
+            # Tool references are rendered by name; emit as text so the
+            # renderer downstream is content-format agnostic.
+            blocks.append(
+                {
+                    "type": MelodyContentType.TEXT,
+                    "text": part.get("name") or part.get("text", ""),
+                }
+            )
+        else:
+            # Unknown block type: render as text fallback.
+            text = part.get("text") or part.get(part_type) or ""
+            text_str = text if isinstance(text, str) else json.dumps(text)
+            blocks.append({"type": MelodyContentType.TEXT, "text": text_str})
+    return blocks
+
+
+def _document_to_melody(doc: Any) -> dict[str, Any]:
+    """Coerce a Cohere v2 document into melody's ``Document`` (dict) shape."""
+    match doc:
+        case str():
+            return {"text": doc}
+        # Cohere v2 wraps documents in {id, data: {...}}; the ``**rest``
+        # capture keeps the outer ``id`` (and any other outer keys) so we
+        # can lift them onto the payload without losing them.
+        case {"data": dict() as data, **rest}:
+            payload = dict(data)
+            if (outer_id := rest.get("id")) is not None and "id" not in payload:
+                payload["id"] = outer_id
+            return payload
+        case dict():
+            return dict(doc)
+        case _:
+            raise TypeError(f"Unsupported document type: {type(doc).__name__}")
+
+
+def _tool_to_melody(tool: Any) -> dict[str, Any]:
+    """Coerce a chat completions tool definition into melody's ``Tool`` shape.
+
+    Accepts either the raw OpenAI tool wrapper ``{type:"function", function:
+    {name, description, parameters}}`` or a flat ``{name, description,
+    parameters}`` dict (which is what melody itself expects).
+    """
+    if hasattr(tool, "model_dump"):
+        tool = tool.model_dump()
+    if not isinstance(tool, dict):
+        raise TypeError(f"Unsupported tool type: {type(tool).__name__}")
+    if "function" in tool and isinstance(tool["function"], dict):
+        fn = tool["function"]
+    else:
+        fn = tool
+    return {
+        "name": fn.get("name", ""),
+        "description": fn.get("description", "") or "",
+        "parameters": fn.get("parameters") or {},
+    }
+
+
+def _conversation_to_melody_messages(
+    conversation: list[ConversationMessage],
+    messages_citations: dict[int, list[dict[str, Any]]] | None = None,
+) -> list[dict[str, Any]]:
+    """Convert vLLM ``ConversationMessage``s into melody's message dict shape.
+
+    ``messages_citations`` is an optional index-keyed lookup of
+    per-message citations from the request, populated by
+    ``CohereServingChatV2`` under the
+    :data:`MESSAGES_CITATIONS_KEY` chat_template_kwargs entry. Values
+    must already be in melody's ``FilterCitation`` dict shape (see
+    ``CohereServingChatV2._sdk_citation_to_melody``). Index is the
+    position in the *input* ``conversation`` list, which currently maps
+    1-to-1 with ``CohereChatV2Request.messages`` for text-only Cohere
+    requests. This mapping breaks silently if ``parse_chat_messages``
+    ever expands a single input message into multiple entries (e.g.
+    multi-modal splitting), which is not the case today.
+    """
+    out: list[dict[str, Any]] = []
+    for i, msg in enumerate(conversation):
+        role = _role_to_melody(msg.get("role", "user"))
+        content_blocks = _content_blocks(msg.get("content"))
+
+        # Treat reasoning content as a thinking block on assistant turns
+        # so multi-turn reasoning is preserved in the rendered prompt.
+        #
+        # Cohere models output thought as either a ``thinking``
+        # block (reasoning models) or a ``tool_plan`` field (older non-
+        # reasoning Command models), but vLLM's ConversationMessage has a
+        # single unified ``reasoning`` field that drops that difference.
+        # The cmd3 / cmd4 jinja templates render ``thinking`` blocks in both
+        # the case of tool calls and regular thinking blocks, so this is fine
+        # from the renderer's perspective.
+        reasoning = msg.get("reasoning") or msg.get("reasoning_content")
+        if role == "chatbot" and reasoning:
+            content_blocks.insert(
+                0,
+                {"type": MelodyContentType.THINKING, "thinking": reasoning},
+            )
+
+        tool_calls = [_normalize_tool_call(tc) for tc in (msg.get("tool_calls") or [])]
+
+        out_msg: dict[str, Any] = {
+            "role": role,
+            "content": content_blocks,
+            "tool_calls": tool_calls,
+        }
+        tool_call_id = msg.get("tool_call_id")
+        if tool_call_id:
+            out_msg["tool_call_id"] = tool_call_id
+        if messages_citations and (cites := messages_citations.get(i)):
+            out_msg["citations"] = cites
+        out.append(out_msg)
+    return out
+
+
+def _build_render_config(
+    conversation: list[ConversationMessage],
+    chat_template_kwargs: dict[str, Any],
+    chat_template: str | None = None,
+) -> tuple[str, dict[str, Any]]:
+    """Build the ``render_cmd3`` / ``render_cmd4`` config dict.
+
+    ``chat_template`` is the standard vLLM per-request template override
+    (from ``ChatCompletionRequest.chat_template`` / the server-side
+    ``--chat-template`` flag). When non-``None`` it is forwarded to melody
+    as ``template_jinja`` so a single vLLM-native path drives template
+    source through the trust guard.
+
+    Returns ``(format, config_dict)`` where ``format`` is either ``"cmd3"``
+    or ``"cmd4"``.
+    """
+    fmt = chat_template_kwargs.get("cohere_format", _DEFAULT_FORMAT)
+    if fmt not in _VALID_FORMATS:
+        raise ValueError(
+            f"Invalid cohere_format={fmt!r}; expected one of {_VALID_FORMATS}"
+        )
+
+    # Reject Cohere-only inlets for raw Jinja source. These keys are
+    # valid at Cohere's own API surface but not at vLLM's: callers must
+    # use the standard vLLM ``chat_template`` request field so the
+    # ``--trust-request-chat-template`` guard applies uniformly.
+    for cohere_only_key in ("template_jinja", "template"):
+        if chat_template_kwargs.get(cohere_only_key) is not None:
+            raise ValueError(
+                f"chat_template_kwargs.{cohere_only_key!r} is not accepted "
+                "in vLLM (it is a Cohere-only field); pass the template "
+                "body via the standard 'chat_template' request field "
+                "instead (requires --trust-request-chat-template)."
+            )
+
+    config: dict[str, Any] = {
+        "messages": _conversation_to_melody_messages(
+            conversation,
+            chat_template_kwargs.get(MESSAGES_CITATIONS_KEY),
+        ),
+    }
+
+    # ``template_id`` is a selector for one of melody's built-in template
+    # variants (not raw source), so it is safe to accept from the client.
+    if template_id := chat_template_kwargs.get("template_id"):
+        config["template_id"] = template_id
+    # Raw template source flows exclusively through the standard vLLM
+    # ``chat_template`` field so it is subject to the
+    # ``--trust-request-chat-template`` guard in OnlineRenderer.
+    if chat_template:
+        config["template_jinja"] = chat_template
+    # Only support Jinja with vllm
+    config["use_jinja"] = True
+
+    # Documents
+    documents = chat_template_kwargs.get("documents") or []
+    if documents:
+        config["documents"] = [_document_to_melody(d) for d in documents]
+
+    # Tools - prefer explicit ``available_tools``, fall back to OpenAI ``tools``
+    tools = (
+        chat_template_kwargs.get("available_tools")
+        or chat_template_kwargs.get("tools")
+        or []
+    )
+    if tools:
+        config["available_tools"] = [_tool_to_melody(t) for t in tools]
+
+    # Reasoning toggle (cmd3 + cmd4)
+    if (rt := chat_template_kwargs.get("reasoning_type")) is not None:
+        config["reasoning_type"] = str(rt)
+    elif "thinking" in chat_template_kwargs:
+        # Cohere v2 ``thinking: {type: enabled|disabled}`` shorthand
+        thinking = chat_template_kwargs["thinking"]
+        t = thinking.get("type") if isinstance(thinking, dict) else thinking
+        if t in ("enabled", "disabled"):
+            config["reasoning_type"] = t
+
+    if (di := chat_template_kwargs.get("dev_instruction")) is not None:
+        config["dev_instruction"] = str(di)
+
+    # JSON / structured outputs
+    if (rf := chat_template_kwargs.get("response_format")) is not None:
+        rf = rf.model_dump() if hasattr(rf, "model_dump") else dict(rf)
+        rf_type = rf.get("type")
+        if rf_type == "json_object":
+            config["json_mode"] = True
+        elif rf_type in ("json_schema", "json"):
+            schema = rf.get("schema") or rf.get("json_schema")
+            if isinstance(schema, dict) and "schema" in schema:
+                schema = schema["schema"]
+            if schema is not None:
+                config["json_schema"] = (
+                    schema if isinstance(schema, str) else json.dumps(schema)
+                )
+    if (js := chat_template_kwargs.get("json_schema")) is not None:
+        config["json_schema"] = js if isinstance(js, str) else json.dumps(js)
+    if "json_mode" in chat_template_kwargs:
+        config["json_mode"] = bool(chat_template_kwargs["json_mode"])
+
+    if fmt == "cmd3":
+        if (sm := chat_template_kwargs.get("safety_mode")) is not None:
+            config["safety_mode"] = str(sm).lower()
+        # citation_quality: ``on`` / ``off``
+        cq = chat_template_kwargs.get("citation_quality")
+        if cq is None and (co := chat_template_kwargs.get("citation_options")):
+            mode = co.get("mode") if isinstance(co, dict) else None
+            if mode is not None:
+                cq = "on" if str(mode).lower() != "off" else "off"
+        if cq is not None:
+            config["citation_quality"] = str(cq).lower()
+        if "skip_preamble" in chat_template_kwargs:
+            config["skip_preamble"] = bool(chat_template_kwargs["skip_preamble"])
+    else:  # cmd4
+        # cmd4 uses ``grounding`` rather than safety_mode/citation_quality.
+        # melody's cmd4 only accepts ``unknown`` / ``enabled`` / ``disabled``,
+        # so the Cohere v2 ``citation_options.mode`` values
+        # (``FAST`` / ``ACCURATE`` / ``OFF``) have to be normalized -- a
+        # raw lowercased passthrough would raise from
+        # ``render_cmd4`` (cmd4 has no fast/accurate distinction at the
+        # prompt-template layer; both request grounding-on).
+        if (gr := chat_template_kwargs.get("grounding")) is not None:
+            config["grounding"] = _normalize_cmd4_grounding(gr)
+        elif co := chat_template_kwargs.get("citation_options"):
+            mode = co.get("mode") if isinstance(co, dict) else None
+            if mode is not None:
+                config["grounding"] = _normalize_cmd4_grounding(mode)
+        if (pi := chat_template_kwargs.get("platform_instruction")) is not None:
+            config["platform_instruction"] = str(pi)
+
+    # Anything we didn't explicitly interpret above is forwarded to melody
+    # as a Jinja template variable. This matches vLLM's documented contract
+    # for ``chat_template_kwargs`` ("kwargs accessible by the template")
+    # and lets callers write e.g. ``{"reasoning_effort": "low"}`` directly
+    # without a nested ``additional_template_fields`` wrapper.
+    extra = {
+        k: v
+        for k, v in chat_template_kwargs.items()
+        if k not in _RENDERER_CONSUMED_KEYS
+    }
+    if extra:
+        config["additional_template_fields"] = extra
+
+    return fmt, config
+
+
+class CohereRenderer(BaseRenderer[HfTokenizer]):
+    """Renderer that templates Cohere prompts via the melody Rust bindings.
+
+    Tokenization is delegated to the standard HF tokenizer; only the
+    chat-template step is replaced with ``cohere_melody.render_cmd3`` /
+    ``render_cmd4``. Enabled via ``--tokenizer-mode cohere``.
+    """
+
+    def __init__(
+        self,
+        config: VllmConfig,
+        tokenizer: HfTokenizer | None,
+    ) -> None:
+        # Match HfRenderer in not mutating the cached tokenizer instance
+        tokenizer = copy.copy(tokenizer)
+        super().__init__(config, tokenizer)
+
+        # Lazy import to keep `cohere_melody` an optional dependency
+        self._melody = _try_import_melody()
+        # ``render_cmd3`` / ``render_cmd4`` are pure CPU work; cache the
+        # thread-pool wrapper once so the async path doesn't allocate a new
+        # adapter on every request.
+        self._render_async = make_async(self._render, executor=self._executor)
+
+    def _render(self, fmt: str, config_dict: dict[str, Any]) -> str:
+        if fmt == "cmd3":
+            return self._melody.render_cmd3(config_dict)
+        return self._melody.render_cmd4(config_dict)
+
+    def render_messages(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        params: ChatParams,
+    ) -> tuple[list[ConversationMessage], DictPrompt]:
+        conversation, mm_data, mm_uuids = parse_chat_messages(
+            messages,
+            self.model_config,
+            content_format="openai",
+            media_io_kwargs=params.media_io_kwargs,
+            mm_processor_kwargs=params.mm_processor_kwargs,
+        )
+
+        chat_template_kwargs = dict(params.chat_template_kwargs)
+        fmt, config_dict = _build_render_config(
+            conversation, chat_template_kwargs, params.chat_template
+        )
+        prompt_text = self._render(fmt, config_dict)
+        prompt = parse_dec_only_prompt(prompt_text)
+
+        if mm_data is not None:
+            prompt["multi_modal_data"] = mm_data
+        if mm_uuids is not None:
+            prompt["multi_modal_uuids"] = mm_uuids
+
+        return conversation, prompt
+
+    async def render_messages_async(
+        self,
+        messages: list[ChatCompletionMessageParam],
+        params: ChatParams,
+    ) -> tuple[list[ConversationMessage], DictPrompt]:
+        conversation, mm_data, mm_uuids = await parse_chat_messages_async(
+            messages,
+            self.model_config,
+            content_format="openai",
+            media_io_kwargs=params.media_io_kwargs,
+            mm_processor_kwargs=params.mm_processor_kwargs,
+        )
+
+        chat_template_kwargs = dict(params.chat_template_kwargs)
+        fmt, config_dict = _build_render_config(
+            conversation, chat_template_kwargs, params.chat_template
+        )
+        prompt_text = await self._render_async(fmt, config_dict)
+        prompt = parse_dec_only_prompt(prompt_text)
+
+        if mm_data is not None:
+            prompt["multi_modal_data"] = mm_data
+        if mm_uuids is not None:
+            prompt["multi_modal_uuids"] = mm_uuids
+
+        return conversation, prompt

vllm/renderers/registry.py

@@ -20,6 +20,7 @@ logger = init_logger(__name__)
 
 
 _VLLM_RENDERERS = {
+    "cohere": ("cohere", "CohereRenderer"),
     "deepseek_v32": ("deepseek_v32", "DeepseekV32Renderer"),
     "deepseek_v4": ("deepseek_v4", "DeepseekV4Renderer"),
     "hf": ("hf", "HfRenderer"),

vllm/tokenizers/registry.py

@@ -40,6 +40,9 @@ _MODEL_TYPES_WITH_INCORRECT_TOKENIZER_CLASS: set[str] = {
 }
 
 _VLLM_TOKENIZERS = {
+    # ``cohere`` mode uses the standard cached HF tokenizer; only the
+    # renderer (template stage) is replaced with a melody-based one.
+    "cohere": ("hf", "CachedHfTokenizer"),
     "deepseek_v32": ("deepseek_v32", "DeepseekV32Tokenizer"),
     "deepseek_v4": ("deepseek_v4", "DeepseekV4Tokenizer"),
     "hf": ("hf", "CachedHfTokenizer"),