feat: reusable rpc refs

[mhuisi]

This PR adds support for server-sided RpcRef reuse and fixes a bug where trace nodes in the InfoView would close while the file was still being processed.

The core of the trace node issue is that the server always serves new RPC references in every single response to the client, which means that the client is forced to reset its UI state.

In a previous attempt at fixing this (#8056), the server would memorize the RPC-encoded JSON of interactive diagnostics (which includes RPC references) and serve it for as long as it could reuse the snapshot containing the diagnostics, so that RPC references are reused. The problem with this was that the client then had multiple finalizers registered for the same RPC reference (one for every reused RPC reference that was served), and once the first reference was garbage-collected, all other reused references would point into the void.

This PR takes a different approach to resolve the issue: The meaning of $/lean/rpc/release is relaxed from "Free the object pointed to by this RPC reference" to "Decrement the RPC reference count of the object pointed to by this RPC reference", and the server now maintains a reference count to track how often a given RpcRef was served. Only when every single served instance of the RpcRef has been released, the object is freed. Additionally, the reuse mechanism is generalized from being only supported for interactive diagnostics, to being supported for any object using WithRpcRef. In order to make use of reusable RPC references, downstream users still need to memorize the WithRpcRef instances accordingly.

Closes #8053. This PR should only be merged at the start of the next release cycle.

Breaking changes

Since WithRpcRef is now capable of tracking its identity to decide which WithRpcRef usage constitutes a reuse, the constructor of WithRpcRef has been made private to discourage downstream users from creating WithRpcRef instances with manually-set ids. Instead, WithRpcRef.mk (which lives in BaseIO) is now the preferred way to create WithRpcRef instances.

[leanprover-community-bot]

Mathlib CI status (docs):

• ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 8195f7050283953d361e221d74c7607f6396b0d7 --onto 416e07a68e5d9b884de8af4836497a7cc2414c8f. You can force Mathlib CI using the force-mathlib-ci label. (2025-04-25 16:31:39)
• ❗ Batteries/Mathlib CI will not be attempted unless your PR branches off the nightly-with-mathlib branch. Try git rebase 2594a8edadd163ba579fc4a7b5a989650f65edda --onto 4eccb5b4792c270ad10ac059b9672a8845961079. You can force Mathlib CI using the force-mathlib-ci label. (2025-05-22 17:11:01)

[Vtec234]

The approach makes sense, but it feels very complicated for what it does. I am wondering about alternatives.

• First of all, I am assuming that there is (still) no way for the server to only send each diagnostic once, when its snapshot first finishes computing? Maybe because the LSP client will clear all other diagnostics upon receiving publishDiagnostics, so the Lean server has to keep resending them?
• What if instead of implementing this server-side, we tagged each Lean.Lsp.Diagnostic with a unique ID, and have the infoview only call getInteractiveDiagnostics on new diagnostics that it receives, caching the results for diagnostic IDs it already knows? It seems that would decrease the amount of data transferred between client and server, and also absolve the server from having to cache anything.
• Supposing the above questions have definite 'no' answers, I have some concerns about the implementation
  • WithRpcRef was intended as a type marker used in client code that derives RpcEncodable. There is by now a good amount of widgets that use WithRpcRef.mk, and privateizing it would break that code. Atm I can't think of a usecase for WithRpcRef.mkReusable outside of this issue that's very specific to diagnostics; I guess widget messages could also use mkReusable in order to maintain their state while new diagnostics stream in? But this requires changing their implementation, while something like the infoview change I suggest above may get away without necessitating that?
  • The global freshWithRpcRefId tracker, combined with the fact that we now need two address spaces - one for IDs and one for RpcRefs - seems quite hacky. Could mkReusable live in StateM RpcObjectStore and just bump nextRef?

[mhuisi]

Thanks for the review! :-)

First of all, I am assuming that there is (still) no way for the server to only send each diagnostic once, when its snapshot first finishes computing? Maybe because the LSP client will clear all other diagnostics upon receiving publishDiagnostics, so the Lean server has to keep resending them?

Yes, exactly. publishDiagnostics does not support incrementally streaming data in this manner.

What if instead of implementing this server-side, we tagged each Lean.Lsp.Diagnostic with a unique ID, and have the infoview only call getInteractiveDiagnostics on new diagnostics that it receives, caching the results for diagnostic IDs it already knows? It seems that would decrease the amount of data transferred between client and server, and also absolve the server from having to cache anything.

This would help with the trace node issue, but it won't help any other users of the RPC protocol. FWIW, I think that it would also be a bit more complicated than simply only calling getInteractiveDiagnostics on new diagnostics - e.g. when you move the cursor, I'd expect the UI state to persist if your cursor is still in the same diagnostic.
From a client perspective, it also seems a little strange that I can't request the same information in the exact same context twice without receiving different responses, so implementing a solution at the RPC level that other users can utilize as well seemed like the right general solution to me.

WithRpcRef was intended as a type marker used in client code that derives RpcEncodable. There is by now a good amount of widgets that use WithRpcRef.mk, and privateizing it would break that code.

If they do really just use WithRpcRef.mk, then the code will continue working as before, since WithRpcRef.mk is now an alias for WithRpcRef.mkNonReusable :-) But yes, if they use any of the constructor syntaxes (also for pattern matching), then their code will break.

Atm I can't think of a usecase for WithRpcRef.mkReusable outside of this issue that's very specific to diagnostics; I guess widget messages could also use mkReusable in order to maintain their state while new diagnostics stream in?

Pretty much any user of the RPC protocol that wants their state to persist for a bit longer could benefit from this, e.g. when moving the cursor but the context hasn't changed or as new diagnostics are still trickling in. IMO fixing this kind of issue for widgets as well at some point in the future is definitely desirable, even if not quite as urgent.
There's of course a question of how they should memorize (and invalidate) their responses. Perhaps we also need to provide a general memorization mechanism at the Snapshot level :-)

The global freshWithRpcRefId tracker, combined with the fact that we now need two address spaces - one for IDs and one for RpcRefs - seems quite hacky. Could mkReusable live in StateM RpcObjectStore and just bump nextRef?

I think it could, but it would complicate client code even more, no? RPC procedures would then need access to the session ID so that they can obtain the RPC session and its RpcObjectStore just to obtain a fresh ID. I suppose nextRef could also be moved out of RpcObjectStore since it doesn't really gain anything from being scoped per-session, though I don't think that it makes a much of a difference. What do you think?

[Vtec234]

FWIW, I think that it would also be a bit more complicated than simply only calling getInteractiveDiagnostics on new diagnostics - e.g. when you move the cursor, I'd expect the UI state to persist if your cursor is still in the same diagnostic.

Oh I forgot that getInteractiveDiagnostics is called whenever the cursor moves around.. why is that, again? Could we not call it exactly once when receiving new publishDiagnostics in WithLspDiagnosticsContext?

From a client perspective, it also seems a little strange that I can't request the same information in the exact same context twice without receiving different responses

It's not so strange if the spec for RPC functions is 'allocate stuff and return a pointer'. Then it's the client's responsibility to cache.

But yes, if they use any of the constructor syntaxes (also for pattern matching), then their code will break.

Yeah sorry, I meant they use anon constructors.

There's of course a question of how they should memorize (and invalidate) their responses.

Yeah this is exactly the problem; for diagnostics the server can cache these WithRpcRefs, but there is no mechanism for general @[rpc_server_method]s to cache anything so they'd have to resort to IO.Ref hacks.

[mhuisi]

Could we not call it exactly once when receiving new publishDiagnostics in WithLspDiagnosticsContext?

I think we could, yes.

Then it's the client's responsibility to cache.

How would this work for widgets? Say I move the cursor around and the widget shouldn't have to re-render while this is happening.
In my mind, the server is the only one that has enough semantic information to decide whether the context that some RPC procedure operates in has changed or not in all cases, and clients can only approximate this in specific cases.
I suppose we might be able to fix the new-diagnostics-reset-state issue in particular for widgets on the client-side by only requesting widgets when the position of the cursor changes? I guess arbitrary RPC requests could depend on diagnosticsRef, but perhaps that should be considered a bug in that code.
Maybe it's worth doing both of these things so that widgets have reasonable default behavior, but there's also the possibility to reuse state even when the cursor position has changed :-)

Yeah this is exactly the problem; for diagnostics the server can cache these WithRpcRefs, but there is no mechanism for general @[rpc_server_method]s to cache anything so they'd have to resort to IO.Ref hacks.

We could add support for this at a later point. I think a basic implementation could be as simple as extending the cache in the snapshot to include arbitrary Dynamic data and provide snapshot lookup functions that also cache the response in an opaque manner in the snapshot that it was sourced from, which would adequately resolve the invalidation question.

[Vtec234]

To summarize the discussion so far, my understanding is that:

• Trace nodes automatically close while file is still elaborating #8053 could potentially be fixed by only calling getInteractiveDiagnostics once in the infoview, somewhere in WithLspDiagnosticsContext. As a nice side effect, this would also reduce the amount of client-server traffic.
• Implementing this PR would provide a mechanism for preserving the state of widget UI elements in more situations, but it would break some existing code.

Assuming the above is correct, I am not opposed to this change since you make some good points, but I think it might make sense to hold off on it before a concrete example of a widget that would benefit from this comes up, and in the meantime just fix this in vscode-lean4; I find it difficult to imagine how such widgets would work in a vacuum. What do you think?

[mhuisi]

Just to summarize what the client-sided fix would look like:

• For the client to be able to tell that it doesn't have to re-request the interactive diagnostics for every new set of diagnostics, the server needs to tag each plain diagnostic with an ID so that it can actually identify diagnostics across publishDiagnostics notifications.
• Additionally, the protocol for getInteractiveDiagnostics needs to be adjusted so that it doesn't simply request the diagnostics for the whole file, but only for a specific set of plain diagnostic IDs. The interactive diagnostics for an ID that has already been requested need to be memorized by the client, which ensures that we don't get new RPC references for previously-requested interactive diagnostics. This will break existing clients.
• If we additionally want the interactive diagnostics to remain stable even when the cursor is moved, we also need to adjust the client to only request the interactive diagnostics when new plain diagnostics are received, not when the cursor is moved.

To be honest, all of this doesn't really seem less painful than the changes in this PR. It's also worth pointing out that NVim has the same problem as VS Code here, and would also have to make all of these adjustments.

Generally, I think I'd prefer to merge this PR sooner, rather than later. It will only get more difficult to make breaking changes in the RPC protocol as time goes on, and I think our conversation here has established that there are cases where the client can't possibly tell whether its information is stale or not and it must ask the server, so we might need to make this change eventually regardless.

[Vtec234]

The interactive diagnostics for an ID that has already been requested need to be memorized by the client, which ensures that we don't get new RPC references for previously-requested interactive diagnostics. This will break existing clients.

Making this change in some clients wouldn't make other clients any more broken than they already are, right?

To be honest, all of this doesn't really seem less painful than the changes in this PR. It's also worth pointing out that NVim has the same problem as VS Code here, and would also have to make all of these adjustments.

I think you are right about both of these, but if the two solutions really are roughly equivalent in complexity, I would still advocate for the mostly-client-based one. It modifies only implementation details; the present solution changes the APIs that Lean users work with without being demonstrably more useful for those users: we do not have any examples of widgets that would make use of reusable RPC refs. I would feel more positive about this if there were such an example. What I am specifically concerned about is that we make this breaking change in order to fix the diagnostics resetting, then realize that it is not actually useful for other widgets, and have to make more breaking changes or hack around the design.

[mhuisi]

What I am specifically concerned about is that we make this breaking change in order to fix the diagnostics resetting, then realize that it is not actually useful for other widgets, and have to make more breaking changes or hack around the design.

The surface area of the breaking change in this PR is very small. All it does is make a constructor private. We could as well eliminate most of the breaking change by just not making it private at all, setting id? := none by default and trusting users not to set WithRpcRef.id? themselves. If you'd prefer that to avoid most of this breaking change, then I'm happy to do that.
Existing widget code would continue using non-reusable RPC refs by default and clients that use a mechanism similar to vscode-lean4 for releasing references will continue working as-is. The only code that should break from this PR is code that uses anonymous constructor notation, which is inherently fragile.

It's perhaps also worth pointing out again that the alternative is not actually a client-only fix. We need to adjust the protocol of publishDiagnostics and the protocol of getInteractiveDiagnostics. The former is a somewhat backwards-compatible extension, whereas for the latter we would likely have to add another RPC method getInteractiveDiagnostics' with diagnostic ID compatible semantics so that clients from before the change can still talk to language servers from after the change.

we do not have any examples of widgets that would make use of reusable RPC refs

It's not difficult to imagine how most widgets with any form of UI state might benefit from them: They allow widgets to retain state when the text cursor is moved in the scope where the same widget is triggered.

In general, I feel that this discussion about using diagnostic IDs vs. using RPC refs is a bit of a red herring: There's good (orthogonal) reasons to have both, because clients generally can't decide whether to reuse state on their own, except in special cases.

• State reuse for plain diagnostics needs diagnostic IDs and cannot benefit from RPC refs (e.g. some editor decorations could benefit from the ability to be able to distinguish freshly issued diagnostics from re-issued ones)
• Reuse for interactive diagnostics can benefit from both (e.g. traces)
• Reuse for other RPC requests can only benefit from RPC refs (e.g. widgets)

To me, the main concern for the alternative proposal is that clients would have to implement several heuristics to request data from the language server less often so that their behavior is more functionally correct. That just doesn't seem like a great API to me - for RPC, the server should have some tool available to tell clients that they may reuse state.

[Vtec234]

Ok, let's zoom out and analyze what is going on here. Why do traces collapse when the cursor moves? Because the infoview makes a new request, the server replies with a new RPC reference, and the new reference causes the UI to reset. But why does the UI reset upon receiving a new reference? Because the entire infoview codebase assumes that object identity (in Lean, Eq or BEq; in JS, "deep equality") implies reference identity (in JS, ===). Then by contraposition, observing a different reference means the data has changed. This is done because object identity checks are expensive, and reference identity checks are cheap; it's the same idea as hash consing. To make this work, the infoview codebase takes pains to ensure that JS object references only change over time when the underlying data has changed (example), and conversely that references do change when the stored data has changed (this can be violated with mutable store).

With the above in mind, what we seem to have disagreed on so far is whether it should be the JS client's or the Lean server's obligation to ensure that object identity for Lean objects implies RPC reference identity (the converse is true because we always use fresh reference IDs). When put like this, it is obvious that the Lean side should do this, as only it knows the underlying Lean objects. (So I agree with you 🙂)

While I think we agree on the above, I had concerns about the implementation. WithRpcRef was intended as a marker that users put on fields to indicate how they should be serialized. I worry that adding the (non)reusable choice will make an already-confusing API even more confusing for users.

I am now starting to think that what's at fault is the decision to store the RPC pointer in WithRpcRef rather than deferring its creation to RpcObjectStore. Here is a third proposal that might solve the problem server-side without adding user-observable complexity.

Let's go back to the identity implications. Is it ever useful for JS to receive different RPC references to the same Lean object across multiple RPC calls? I claim the answer is no (at least, I cannot think of any reasonable code that would rely on this). So what we can do is make all RPC references reusable, but in the sense that their Lean pointer identity implies RPC reference identity. I think this could be done by adjusting Server.rpcStoreRef to look at the pointer to the underlying data of any : Dynamic, and storing that in a map analogous to refsById. It might have to be unsafe, or maybe one can use withPtrAddr.

This means the Lean code that produces WithRpcRefs has to take care to keep its Lean objects around when it wants them cached on the JS side, but it is transparent, non-API-breaking, and seems to make sense conceptually. It also looks like it'd just work for diagnostics as those are already preserved in diagnosticsRef; but I'd prefer this solution even if the server has to do more work to cache diagnostic objects. What do you think?

(A fourth proposal would be for the server to cache references based on Lean object identity as determined by BEq or DecidableEq or similar, but this can be more expensive so should not be a default.)

[Vtec234]

Here is a quick prototype. It seems to fix the trace issue, but I didn't test it very much so there may certainly be something very wrong.

Since we now create a unique RPC ref for each Lean ptr, I cut out the middleman and just made the refs take on Lean pointer values (which was also how the very first RPC prototype worked, if you remember). I think the comment about 'avoiding any possible bugs resulting from RPC ref reuse' was confused; any such bug is a real bug in the client that should be fixed. Unfortunately it looks like the RpcEncodable (WithRpcRef α) instance must be unsafe again; it seems type-safe, but it will produce different JSON on definitionally equal values that are not pointer-equal.

[mhuisi]

Perhaps I should have mentioned this in the PR header: I had more or less this exact design you propose at some point before submitting this PR for similar reasons and rejected it before arriving at the current approach. The main problems are that there is (AFAICT) no way to state in Lean code that two equal values should have different references (e.g. if CSE gets in the way) if you ever had to and that ptrEq a b doesn't necessarily imply a = b at the moment (#1502 - I think that all known examples rely on erasure and that Cameron has plans to fix this at some point in the new codegen, but it definitely doesn't instill confidence in me that we should expose it in our RPC protocol). I also think that it is pretty confusing.

Is it ever useful for JS to receive different RPC references to the same Lean object across multiple RPC calls? I claim the answer is no (at least, I cannot think of any reasonable code that would rely on this).

I'm not sure about this, but I don't have any concrete counter-examples either. With the pointer-based approach, you could definitely even issue the same RPC reference twice in the same response by accident, which may or may not complicate things on the client-side (e.g. if the client-side were to assume that the UI state of two separate occurrences of the same RPC reference can be reused). I'd prefer to err on the side of caution here.

This means the Lean code that produces WithRpcRefs has to take care to keep its Lean objects around when it wants them cached on the JS side, but it is transparent, non-API-breaking, and seems to make sense conceptually.

I don't think that the transparency is an advantage here. In general, when I use Lean, I don't ever need to worry about whether a value I created is referentially identical to another one or just structurally identical. I can cache the value, or I can recreate it with the same value - it won't make a functional difference - and I can always simply look at the definition of the value to figure out how to work with it. When using the pointer of a value to decide which RPC reference it should map to, this is no longer true: In order to reuse an RPC reference, I must now make sure that the value is memorized (not "cached", since it affects the functional behavior), and that I don't recreate the same value somewhere by accident.
The WithRpcRef approach in this PR makes all of this very explicit. I can look at the definition and figure out what to do with the value. I don't need to ensure that I never recreate another value that is only structurally identical by accident. I can directly control how and when RPC references are reused. It's all there, in Lean code, and a correct compiler is sure to respect it.

[Vtec234]

ptrEq a b doesn't necessarily imply a = b at the moment

Oh no! I was not aware of this. FWIW these counterexamples shouldn't matter in practice - they are erased and don't affect runtime behaviour by definition - but..

The WithRpcRef approach in this PR makes all of this very explicit. [..] I can directly control how and when RPC references are reused.

..I see your point here. Reasoning about pointer equality is unavoidable in JS because that's how React compares deps, but at least your design makes the identification of references explicit and easy to manipulate on the Lean side.

discussion about using diagnostic IDs vs. using RPC refs is a bit of a red herring: There's good (orthogonal) reasons to have both

They are not entirely orthogonal: they do different things, but overlap in solving #8053, which is what I understood this PR to be mostly aimed at. Having thought about it more, there are two designs I would be happy with. Below, I respond to some of your concerns about either.

1. Diagnostic IDs, as we discussed.

   adjust the protocol of publishDiagnostics [..] is a somewhat backwards-compatible extension

   Is it not entirely backwards-compatible? We'd be adding a new field, and clients can happily ignore it.

   have to add another RPC method getInteractiveDiagnostics' with diagnostic ID compatible semantics

   I don't think so. GetInteractiveDiagnosticsParams could get a new ids? field which, if present, means the server should respond with interactive versions of only those diagnostics, and if absent would default to the current behaviour. I guess it should be an error to specify both ids? and lineRange?.

2. This PR, but we make mkReusable the default (might as well call it mk then), make WithRpcRef.id? non-optional, and get rid of the pure mk.

   I remain concerned about making the API more complex by introducing a choice whose use will not be clear until you have written a bunch of widgets and observed problems with the UI state resetting. Replacing mk entirely would alleviate this.

   The disadvantage is that computations must now live in BaseIO, but all RPC stuff happens in RequestM or some ElabM anyway, so I don't think this is an issue.

   (e.g. if the client-side were to assume that the UI state of two separate occurrences of the same RPC reference can be reused). I'd prefer to err on the side of caution here.

   I'd rather not build defensive server-side APIs in anticipation of hypothetical client bugs; these should be fixed in clients. (I am of course guilty of doing this, e.g. with the 'avoid any possible bugs' comment.) Anyhow, even if mkReusable is the default, users may choose to discard the WithRpcRef if they want a fresh RPC reference.

   the client can't possibly tell whether its information is stale or not and it must ask the server, so we might need to make this change eventually regardless.
   Reuse for other RPC requests can only benefit from RPC refs (e.g. widgets)

   Maybe this is not what you are saying, but just to be clear, no RPC ref reuse mechanism is necessary for widgets to preserve their UI state. They can just use local ID pools, or other mechanisms appropriate for their specific implementation. Here is an example that avoids calling printMsg when the cursor moves between line numbers with the same parity. The only thing that is actually missing from the core APIs is a way to put data in Snapshots, as you highlighted.

I am still partial to just doing 1., but I would be happy with the proposed variant of 2 (or both, if 1. would be helpful for editor decorations and other things). Please let me know what you think.

[mhuisi]

Let's go with the second approach then - keeping the API simpler in return for breaking a bit more code now is a reasonable compromise. I definitely also agree that we may want to have both eventually (as well as additional heuristics in vscode-lean4 to request information less often) when it becomes more urgent. All of your other points are good as well.

[Vtec234]

Thanks a lot for bearing with me and this long discussion! Just one remaining question and some docstring suggestions.

[Vtec234]

Suggested change

	Value to use for the next fresh `RpcRef`. It is monotonically increasing to avoid any possible
	Value to use for the next fresh `RpcRef`, monotonically increasing.

Per discussion, I think my comment was misguided.

[Vtec234]

Suggested change

	Two `WithRpcRef`s with the same `id` will yield the same `RpcRef`.
	Two `WithRpcRef`s with the same `id` will yield the same client-side reference.

Maybe we can treat RpcRef as an impl detail and avoid reference in user-facing docs.

[Vtec234]

Suggested change

	id : USize
	private id : USize

What do you think about making the projection private as well? It feels like other modules should not be able to observe this.

[Vtec234]

Suggested change

	Two `WithRpcRef`s with the same `id` will yield the same `RpcRef`.
	Hence, storing a `WithRpcRef` produced by `WithRpcRef.mk` and serving it to the client twice will
	also yield the same `RpcRef` twice, allowing clients to reuse their associated UI state across
	RPC requests.
	As long as the client holds at least one reference to this `WithRpcRef`,
	serving it again will yield the same client-side reference.
	Thus, when used as React deps,
	client-side references can help preserve UI state across RPC requests.

I tried to clarify that at least RC=1 on the server is necessary to actually serve the same RpcRef, and suggested using as React dep explicitly.