Add workDoneProgress protocol support

lsp_client/__init__.py

@@ -24,6 +24,15 @@
     TextDocumentIdentifier,
     TextDocumentItem,
     TextDocumentPositionParams,
+    WorkDoneProgressBegin,
+    WorkDoneProgressCancelNotification,
+    WorkDoneProgressCancelParams,
+    WorkDoneProgressCreateParams,
+    WorkDoneProgressCreateRequest,
+    WorkDoneProgressEnd,
+    WorkDoneProgressOptions,
+    WorkDoneProgressParams,
+    WorkDoneProgressReport,
     # Backwards-compatible aliases
     TextDocumentDidChangeRequest,
     TextDocumentDidOpenRequest,
@@ -57,6 +66,15 @@
     "TextDocumentIdentifier",
     "TextDocumentItem",
     "TextDocumentPositionParams",
+    "WorkDoneProgressBegin",
+    "WorkDoneProgressCancelNotification",
+    "WorkDoneProgressCancelParams",
+    "WorkDoneProgressCreateParams",
+    "WorkDoneProgressCreateRequest",
+    "WorkDoneProgressEnd",
+    "WorkDoneProgressOptions",
+    "WorkDoneProgressParams",
+    "WorkDoneProgressReport",
     # Backwards-compatible aliases
     "TextDocumentDidOpenRequest",
     "TextDocumentDidChangeRequest",

lsp_client/protocol.py

@@ -5,7 +5,7 @@
 for reference, and what a correct and complete implementation should look like.
 """
 
-from typing import Any, List, Optional
+from typing import Any, List, Literal, Optional
 
 from pydantic import BaseModel, Field
 
@@ -265,8 +265,13 @@ def __init__(self, **kwargs: Any) -> None:
         super(CancelRequest, self).__init__(**kwargs)
 
 
+# A progress token is either an integer or a string, supplied by whichever
+# side initiates the progress sequence.
+ProgressToken = int | str
+
+
 class ProgressParams(BaseModel):
-    token: int | str
+    token: ProgressToken
     value: dict
 
 
@@ -276,6 +281,73 @@ def __init__(self, **kwargs: Any) -> None:
         super(ProgressNotification, self).__init__(**kwargs)
 
 
+# Work Done Progress
+# See https://microsoft.github.io/language-server-protocol/specifications/lsp/3.17/specification/#workDoneProgress  # noqa: E501
+
+
+class WorkDoneProgressBegin(BaseModel):
+    """Payload signalling the start of a work done progress sequence.
+
+    Sent as the ``value`` of a ``$/progress`` notification.
+    """
+
+    kind: Literal["begin"] = "begin"
+    title: str
+    cancellable: bool | None = None
+    message: str | None = None
+    percentage: int | None = None
+
+
+class WorkDoneProgressReport(BaseModel):
+    """Payload reporting progress within an ongoing work done sequence."""
+
+    kind: Literal["report"] = "report"
+    cancellable: bool | None = None
+    message: str | None = None
+    percentage: int | None = None
+
+
+class WorkDoneProgressEnd(BaseModel):
+    """Payload signalling the end of a work done progress sequence."""
+
+    kind: Literal["end"] = "end"
+    message: str | None = None
+
+
+class WorkDoneProgressOptions(BaseModel):
+    """Server capability marker for features that support work done progress."""
+
+    workDoneProgress: bool | None = None
+
+
+class WorkDoneProgressCreateParams(BaseModel):
+    token: ProgressToken
+
+
+class WorkDoneProgressCreateRequest(BaseRequest):
+    """Server -> client request asking the client to create a progress token."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        kwargs["method"] = "window/workDoneProgress/create"
+        if isinstance(kwargs.get("params"), WorkDoneProgressCreateParams):
+            kwargs["params"] = kwargs["params"].model_dump(exclude_none=True)
+        super(WorkDoneProgressCreateRequest, self).__init__(**kwargs)
+
+
+class WorkDoneProgressCancelParams(BaseModel):
+    token: ProgressToken
+
+
+class WorkDoneProgressCancelNotification(BaseNotification):
+    """Client -> server notification cancelling a work done progress sequence."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        kwargs["method"] = "window/workDoneProgress/cancel"
+        if isinstance(kwargs.get("params"), WorkDoneProgressCancelParams):
+            kwargs["params"] = kwargs["params"].model_dump(exclude_none=True)
+        super(WorkDoneProgressCancelNotification, self).__init__(**kwargs)
+
+
 # Backwards-compatible aliases for renamed classes
 TextDocumentDidOpenRequest = TextDocumentDidOpenNotification
 TextDocumentDidChangeRequest = TextDocumentDidChangeNotification

tests/test_protocol.py

@@ -18,6 +18,14 @@
     TextDocumentDidOpenNotification,
     TextDocumentIdentifier,
     TextDocumentPositionParams,
+    WorkDoneProgressBegin,
+    WorkDoneProgressCancelNotification,
+    WorkDoneProgressCancelParams,
+    WorkDoneProgressCreateParams,
+    WorkDoneProgressCreateRequest,
+    WorkDoneProgressEnd,
+    WorkDoneProgressOptions,
+    WorkDoneProgressReport,
 )
 
 
@@ -147,3 +155,54 @@ def test_definition_request():
     data = req.model_dump(exclude_none=True)
     assert data["method"] == "textDocument/definition"
     assert data["params"]["textDocument"]["uri"] == "file:///tmp/test.py"
+
+
+def test_work_done_progress_begin():
+    begin = WorkDoneProgressBegin(title="Indexing", percentage=0)
+    data = begin.model_dump(exclude_none=True)
+    assert data == {"kind": "begin", "title": "Indexing", "percentage": 0}
+    # Optional fields with None values are excluded
+    assert "message" not in data
+    assert "cancellable" not in data
+
+
+def test_work_done_progress_report():
+    report = WorkDoneProgressReport(message="halfway", percentage=50)
+    data = report.model_dump(exclude_none=True)
+    assert data == {"kind": "report", "message": "halfway", "percentage": 50}
+
+
+def test_work_done_progress_end():
+    end = WorkDoneProgressEnd(message="done")
+    data = end.model_dump(exclude_none=True)
+    assert data == {"kind": "end", "message": "done"}
+
+
+def test_work_done_progress_end_minimal():
+    data = WorkDoneProgressEnd().model_dump(exclude_none=True)
+    assert data == {"kind": "end"}
+
+
+def test_work_done_progress_options():
+    data = WorkDoneProgressOptions(workDoneProgress=True).model_dump(exclude_none=True)
+    assert data == {"workDoneProgress": True}
+
+
+def test_work_done_progress_create_request():
+    req = WorkDoneProgressCreateRequest(
+        id=1, params=WorkDoneProgressCreateParams(token="token-1")
+    )
+    data = req.model_dump(exclude_none=True)
+    assert data["method"] == "window/workDoneProgress/create"
+    assert data["id"] == 1
+    assert data["params"] == {"token": "token-1"}
+
+
+def test_work_done_progress_cancel_notification():
+    notif = WorkDoneProgressCancelNotification(
+        params=WorkDoneProgressCancelParams(token=42)
+    )
+    data = notif.model_dump(exclude_none=True)
+    assert data["method"] == "window/workDoneProgress/cancel"
+    assert "id" not in data
+    assert data["params"] == {"token": 42}