Skip to content

[Prototype] Add multi-connection IPC support to pyrefly TSP #3199

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
heejaechang wants to merge 1 commit into from
Closed

[Prototype] Add multi-connection IPC support to pyrefly TSP #3199

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 32 additions & 0 deletions Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

107 changes: 100 additions & 7 deletions crates/tsp_types/protocol_generator/tsp.json
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,20 @@
}
],
"requests": [
{
"method": "typeServer/connection",
"typeName": "ConnectionRequest",
"messageDirection": "clientToServer",
"documentation": "Main-connection-only request used to open or close an extra TSP transport. Extra transports must remain read-only and must not be used for LSP traffic.",
"params": {
"kind": "reference",
"name": "ConnectionRequestParams"
},
"result": {
"kind": "reference",
"name": "ConnectionRequestResult"
}
},
{
"method": "typeServer/getComputedType",
"typeName": "GetComputedTypeRequest",
Expand Down Expand Up @@ -301,13 +315,21 @@
"v0_1_0": "0.1.0",
"v0_2_0": "0.2.0",
"v0_3_0": "0.3.0",
"current": "0.4.0"
"v0_4_0": "0.4.0",
"current": "0.5.0"
},
"valueDocumentation": {
"v0_1_0": "Initial protocol version",
"v0_2_0": "Added new request types and fields",
"v0_3_0": "Switch to more complex types",
"current": "Switch to Type union and using stubs"
"v0_4_0": "Switch to Type union and using stubs",
"current": "Add multi-connection negotiation and control requests"
}
},
"ConnectionTransportKind": {
"kind": "stringEnum",
"values": {
"Ipc": "ipc"
}
},
"Variance": {
Expand Down Expand Up @@ -345,6 +367,77 @@
],
"documentation": "Represents a location in source code (a node in the AST). Used to point to specific declarations, expressions, or statements in Python source files. Used for: - Pointing to where a type is declared - Identifying the location of expressions for type inference - Error reporting and diagnostics - Linking types back to their source definitions Examples: - For `def foo():`, the node points to the function declaration - For a variable `x = 42`, the node points to the assignment - For default parameter values in functions"
},
"TypeServerMultiConnectionCapability": {
"kind": "interface",
"properties": [
{
"name": "supportedTransports",
"type": {
"kind": "array",
"element": {
"kind": "reference",
"name": "ConnectionTransportKind"
}
},
"optional": false
}
],
"documentation": "Capability shape exchanged via the LSP initialize request/response under `capabilities.experimental.typeServerMultiConnection`."
},
"ConnectionRequestParams": {
"kind": "interface",
"properties": [
{
"name": "type",
"type": {
"kind": "base",
"name": "string"
},
"optional": false
},
{
"name": "kind",
"type": {
"kind": "reference",
"name": "ConnectionTransportKind"
},
"optional": false
},
{
"name": "args",
"type": {
"kind": "array",
"element": {
"kind": "base",
"name": "string"
}
},
"optional": true
}
],
"documentation": "Main-connection-only control request used to open or close extra read-only TSP channels after LSP initialization has completed."
},
"ConnectionRequestResult": {
"kind": "interface",
"properties": [
{
"name": "success",
"type": {
"kind": "base",
"name": "boolean"
},
"optional": false
},
{
"name": "message",
"type": {
"kind": "base",
"name": "string"
},
"optional": true
}
]
},
"ModuleName": {
"kind": "interface",
"properties": [
Expand Down Expand Up @@ -501,8 +594,8 @@
{
"name": "returnType",
"type": {
"kind": "reference",
"name": "Type"
"kind": "reference",
"name": "Type"
},
"optional": true,
"documentation": "Specialized type of the declared return type. Undefined if there is no declared return type. Example: For `def foo[T](x: T) -> T` specialized to `T=int`, returnType=int."
Expand Down Expand Up @@ -589,7 +682,7 @@
"documentation": "Discriminator field that determines which declaration variant this is. Regular: Has source code and AST node Synthesized: Created by type checker, no source node"
}
],
"documentation": "Base interface for all declaration types. Provides the discriminator field for the Declaration union."
"documentation": "Base interface for all declaration types. Provides the discriminator field for the Declaration union. This is a generic interface that is extended by: - RegularDeclaration (kind = Regular) - SynthesizedDeclaration (kind = Synthesized) The type parameter T ensures that the kind field matches the implementing interface. Used for type-safe discrimination: ```typescript if (declaration.kind === DeclarationKind.Regular) { // TypeScript knows this is RegularDeclaration const node = declaration.node; } ```"
},
"RegularDeclaration": {
"kind": "interface",
Expand All @@ -615,8 +708,8 @@
{
"name": "name",
"type": {
"kind": "base",
"name": "string"
"kind": "base",
"name": "string"
},
"optional": true,
"documentation": "Name of the declared symbol, or undefined for anonymous declarations. Example: \"foo\" for `def foo():`, undefined for lambda functions."
Expand Down
68 changes: 65 additions & 3 deletions crates/tsp_types/src/protocol.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,8 +13,7 @@
// 1. Create tsp.json and tsp.schema.json from typeServerProtocol.ts
// 2. Install lsprotocol generator: `pip install git+https://github.com/microsoft/lsprotocol.git`
// 3. Run: `python generate_protocol.py`
use serde::Deserialize;
use serde::Serialize;
use serde::{Deserialize, Serialize};
use serde_repr::Deserialize_repr;
use serde_repr::Serialize_repr;

Expand Down Expand Up @@ -110,6 +109,8 @@ pub enum LSPNull {

#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
pub enum TSPRequestMethods {
#[serde(rename = "typeServer/connection")]
TypeServerConnection,
#[serde(rename = "typeServer/getComputedType")]
TypeServerGetComputedType,
#[serde(rename = "typeServer/getDeclaredType")]
Expand All @@ -129,6 +130,11 @@ pub enum TSPRequestMethods {
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(tag = "method")]
pub enum TSPRequests {
#[serde(rename = "typeServer/connection")]
ConnectionRequest {
id: serde_json::Value,
params: ConnectionRequestParams,
},
#[serde(rename = "typeServer/getComputedType")]
GetComputedTypeRequest {
id: serde_json::Value,
Expand Down Expand Up @@ -347,9 +353,19 @@ pub enum TypeServerVersion {

/// Switch to Type union and using stubs
#[serde(rename = "0.4.0")]
V040,

/// Add multi-connection negotiation and control requests
#[serde(rename = "0.5.0")]
Current,
}

#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
pub enum ConnectionTransportKind {
#[serde(rename = "ipc")]
Ipc,
}

#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
pub enum Variance {
/// Variance not yet determined, will be inferred
Expand Down Expand Up @@ -418,6 +434,33 @@ pub struct Node {
pub uri: String,
}

/// Capability shape exchanged via the LSP initialize request/response under `capabilities.experimental.typeServerMultiConnection`.
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct TypeServerMultiConnectionCapability {
pub supported_transports: Vec<ConnectionTransportKind>,
}

/// Main-connection-only control request used to open or close extra read-only TSP channels after LSP initialization has completed.
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ConnectionRequestParams {
pub args: Option<Vec<String>>,

pub kind: ConnectionTransportKind,

#[serde(rename = "type")]
pub type_: String,
}

#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ConnectionRequestResult {
pub message: Option<String>,

pub success: bool,
}

/// Represents a Python module name, handling both absolute and relative imports. Used for: - Import statement resolution - Tracking module dependencies - Resolving relative imports (from . import, from .. import) Examples: - `import os.path`: leadingDots=0, nameParts=['os', 'path'] - `from . import utils`: leadingDots=1, nameParts=['utils'] - `from ...parent import module`: leadingDots=3, nameParts=['parent', 'module'] - `import mymodule`: leadingDots=0, nameParts=['mymodule']
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
Expand Down Expand Up @@ -510,7 +553,7 @@ pub struct SentinelLiteral {
pub module_name: String,
}

/// Base interface for all declaration types. Provides the discriminator field for the Declaration union.
/// Base interface for all declaration types. Provides the discriminator field for the Declaration union. This is a generic interface that is extended by: - RegularDeclaration (kind = Regular) - SynthesizedDeclaration (kind = Synthesized) The type parameter T ensures that the kind field matches the implementing interface. Used for type-safe discrimination: ```typescript if (declaration.kind === DeclarationKind.Regular) { // TypeScript knows this is RegularDeclaration const node = declaration.node; } ```
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct DeclarationBase {
Expand Down Expand Up @@ -864,6 +907,25 @@ pub enum LSPIdOptional {
None,
}

/// Main-connection-only request used to open or close an extra TSP transport. Extra transports must remain read-only and must not be used for LSP traffic.
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
pub struct ConnectionRequest {
/// The version of the JSON RPC protocol.
pub jsonrpc: String,

/// The method to be invoked.
pub method: TSPRequestMethods,

/// The request id.
pub id: LSPId,

pub params: ConnectionRequestParams,
}

/// Response to the [ConnectionRequest].
pub type ConnectionResponse = ConnectionRequestResult;

/// Requests and notifications for the type server protocol. Request for the computed type of a declaration or node. Computed type is the type that is inferred based on the code flow. Example: def foo(a: int | str): if instanceof(a, int): b = a + 1 # Computed type of 'b' is 'int'
#[derive(Serialize, Deserialize, PartialEq, Debug, Eq, Clone)]
#[serde(rename_all = "camelCase", deny_unknown_fields)]
Expand Down
2 changes: 1 addition & 1 deletion crates/tsp_types/tests/protocol_types.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ fn test_variance_round_trip() {
fn test_type_server_version_round_trip() {
let v = TypeServerVersion::Current;
let json = serde_json::to_value(&v).unwrap();
assert_eq!(json, serde_json::json!("0.4.0"));
assert_eq!(json, serde_json::json!("0.5.0"));
let back: TypeServerVersion = serde_json::from_value(json).unwrap();
assert_eq!(back, v);
}
Expand Down
1 change: 1 addition & 0 deletions pyrefly/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ faster-hex = "0.6.1"
fuzzy-matcher = "0.3.7"
fxhash = "0.2.1"
indicatif = { version = "0.18.4", features = ["futures", "improved_unicode", "rayon", "tokio"] }
interprocess = "2.2.3"
itertools = "0.14.0"
lsp-server = "0.7.9"
lsp-types = { git = "https://github.com/astral-sh/lsp-types", rev = "3512a9f33eadc5402cfab1b8f7340824c8ca1439" }
Expand Down
8 changes: 5 additions & 3 deletions pyrefly/lib/commands/tsp.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,10 @@ pub struct TspArgs {
/// Note that indexing files is a performance-intensive task.
#[arg(long, default_value_t = if cfg!(fbcode_build) {0} else {2000})]
pub(crate) workspace_indexing_limit: usize,
/// Selects the transport for the main JSON-RPC connection.
/// Use `stdio` (default) or `ipc://<name>` for a local socket / named pipe.
#[arg(long, default_value = "stdio")]
pub(crate) transport: String,
}

pub fn run_tsp(
Expand Down Expand Up @@ -107,9 +111,7 @@ impl TspArgs {
// Note that we must have our logging only write out to stderr.
eprintln!("starting TSP server");

// Create the transport. Includes the stdio (stdin and stdout) versions but this could
// also be implemented to use sockets or HTTP.
let (connection, reader, io_threads) = Connection::stdio();
let (connection, reader, io_threads) = Connection::from_transport(&self.transport)?;

run_tsp(connection, reader, self, telemetry, wrapper, thread_count)?;
io_threads.join()?;
Expand Down
Loading
Loading