Rewrite local association scenario for cleaner error handling and cle…#260
Rewrite local association scenario for cleaner error handling and cle…#260Kuldotha wants to merge 4 commits into
Conversation
|
Warning Rate limit exceeded
Your organization is not enrolled in usage-based pricing. Contact your admin to enable usage-based pricing to continue reviews beyond the rate limit, or try again in 48 minutes and 3 seconds. ⌛ How to resolve this issue?After the wait time has elapsed, a review can be triggered using the We recommend that you space out your commits to avoid hitting the rate limit. 🚦 How do rate limits work?CodeRabbit enforces hourly rate limits for each developer per organization. Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout. Please see our FAQ for further information. ℹ️ Review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughLocalAssociationScenario now implements IDisposable and uses cancellation-driven async StartAndExecute with websocket/backoff/key-exchange handling and centralized cleanup. Three SolanaMobileWalletAdapter call sites use using for LocalAssociationScenario. SessionWallet gains GetSessionWallet factory. Web3.rpcCluster default is now conditional on MAINNET. Changes
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
🚥 Pre-merge checks | ✅ 2 | ❌ 3❌ Failed checks (2 warnings, 1 inconclusive)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Fix all issues with AI agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs`:
- Around line 176-184: ConnectWithBackoffAsync can hang because it awaits
_wsConnected.Task but OnWsError only logs; update the WebSocket error handler
(OnWsError) to signal the waiting task (use _wsConnected.TrySetException with
the exception or _wsConnected.TrySetResult(false)) so ConnectWithBackoffAsync's
await unblocks and the retry/backoff logic proceeds; also ensure the close
handler (OnWsClose/OnClose) uses _wsConnected.TrySetResult(true/false)
consistently and guard against double-sets by using TrySet... methods when
touching _wsConnected to avoid race conditions.
- Around line 53-57: You created a linked CancellationTokenSource with
CancelAfter(_overallTimeout) but still assign and use the original ct, so the
overall timeout never cancels ConnectWithBackoffAsync or the action loop; change
the assignment to use the linked CTS token (set _cancellationToken = cts.Token)
and ensure all cancellation checks and ThrowIfCancellationRequested() calls use
_cancellationToken (and not the original ct); apply the same fix to the other
linked-CTS block (the one around lines with CreateLinkedTokenSource/CancelAfter
at the second site) so the overall timeout is honored across
ConnectWithBackoffAsync and the action loop.
- Around line 138-143: The Debug.Log in StartActivityForAssociation is emitting
sensitive association tokens (and there are similar logs around lines 250-257) —
remove or redact raw tokens/decrypted payloads and instead log only
non-sensitive metadata (e.g., port, token length, a stable prefix/suffix, or a
short hash). Update StartActivityForAssociation and the other logging sites to
stop printing associationToken or decrypted payloads; if a trace is required,
compute and log a one-way hash (SHA-256) or mask all but 4 characters of the
token and include contextual text like "Launched intent for port {port},
token_hash={hash}" or "token_masked=xxxx". Ensure no plaintext tokens or
payloads are written to Debug.Log anywhere in the class.
| using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct); | ||
| cts.CancelAfter(_overallTimeout); | ||
|
|
||
| _cancellationToken = ct; | ||
| _tcs = new TaskCompletionSource<Response<object>>(); |
There was a problem hiding this comment.
Overall timeout isn’t enforced across the connection/action flow.
You create a linked CTS with CancelAfter, but _cancellationToken and ThrowIfCancellationRequested use the original ct, so the 30s timeout won’t cancel ConnectWithBackoffAsync or the action loop when no external token is provided.
✅ Minimal fix to honor the overall timeout
- _cancellationToken = ct;
+ _cancellationToken = cts.Token;
_tcs = new TaskCompletionSource<Response<object>>();
@@
- ct.ThrowIfCancellationRequested();
+ _cancellationToken.ThrowIfCancellationRequested();Also applies to: 94-97
🤖 Prompt for AI Agents
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
53 - 57, You created a linked CancellationTokenSource with
CancelAfter(_overallTimeout) but still assign and use the original ct, so the
overall timeout never cancels ConnectWithBackoffAsync or the action loop; change
the assignment to use the linked CTS token (set _cancellationToken = cts.Token)
and ensure all cancellation checks and ThrowIfCancellationRequested() calls use
_cancellationToken (and not the original ct); apply the same fix to the other
linked-CTS block (the one around lines with CreateLinkedTokenSource/CancelAfter
at the second site) so the overall timeout is honored across
ConnectWithBackoffAsync and the action loop.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Fix all issues with AI agents
In `@Runtime/codebase/SessionWallet.cs`:
- Around line 69-85: GetSessionWallet(string publicKey, string privateKey,
PublicKey targetProgram) lacks the singleton guard and null checks: add the same
early return used in the other overloads (if (Instance != null) return
Instance;) before constructing a new SessionWallet, and validate Web3.Wallet and
its ActiveRpcClient (e.g., if (Web3.Wallet == null) throw/new error or handle;
if (_externalWallet.ActiveRpcClient == null) get node address another way or
throw) to avoid NullReferenceException; ensure you still set _externalWallet =
Web3.Wallet, create sessionAccount, compute SessionTokenPDA via
FindSessionToken, and populate SessionWallet fields only after these checks.
In `@Runtime/codebase/Web3.cs`:
- Around line 31-37: You removed Unity serialization for the RpcCluster field by
adding [NonSerialized] to rpcCluster, which prevents Inspector values from being
saved; restore Inspector configurability by removing [NonSerialized] from the
rpcCluster field (or, if you intend compile-time locking via the MAINNET symbol,
keep the field serialized but make it read-only in the Inspector instead), and
ensure the conditional defaulting using MAINNET still sets the initial value
(rpcCluster and RpcCluster) while preserving serialization or explicitly
document the breaking change if you must keep it non-serialized.
| public static SessionWallet GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram) | ||
| { | ||
| _externalWallet = Web3.Wallet; | ||
|
|
||
| var sessionAccount = new Account(privateKey, publicKey); | ||
|
|
||
| // TODO: ActiveRpcClient can be null, get node address some other way | ||
| var sessionWallet = new SessionWallet(_externalWallet.RpcCluster, _externalWallet.ActiveRpcClient.NodeAddress.ToString()) | ||
| { | ||
| TargetProgram = targetProgram, | ||
| EncryptedKeystoreKey = $"{_externalWallet.Account.PublicKey}_SessionKeyStore", | ||
| SessionTokenPDA = FindSessionToken(targetProgram, sessionAccount, _externalWallet.Account), | ||
| Account = sessionAccount, | ||
| }; | ||
|
|
||
| return sessionWallet; | ||
| } |
There was a problem hiding this comment.
Missing singleton guard — will throw if Instance is already set.
The other GetSessionWallet overloads (lines 97 and 129) check if (Instance != null) return Instance; before constructing a new SessionWallet. This overload skips that check, so calling it when a session wallet already exists will hit the throw new Exception("SessionWallet already exists") in the constructor (line 37).
Additionally, Web3.Wallet (line 71) and _externalWallet.ActiveRpcClient (line 76) can both be null, which would cause a NullReferenceException.
Proposed fix
public static SessionWallet GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram)
{
+ if (Instance != null) return Instance;
+
_externalWallet = Web3.Wallet;
+ if (_externalWallet?.ActiveRpcClient == null)
+ throw new InvalidOperationException("External wallet or its RPC client is not initialized.");
var sessionAccount = new Account(privateKey, publicKey);
- // TODO: ActiveRpcClient can be null, get node address some other way
var sessionWallet = new SessionWallet(_externalWallet.RpcCluster, _externalWallet.ActiveRpcClient.NodeAddress.ToString())🤖 Prompt for AI Agents
In `@Runtime/codebase/SessionWallet.cs` around lines 69 - 85,
GetSessionWallet(string publicKey, string privateKey, PublicKey targetProgram)
lacks the singleton guard and null checks: add the same early return used in the
other overloads (if (Instance != null) return Instance;) before constructing a
new SessionWallet, and validate Web3.Wallet and its ActiveRpcClient (e.g., if
(Web3.Wallet == null) throw/new error or handle; if
(_externalWallet.ActiveRpcClient == null) get node address another way or throw)
to avoid NullReferenceException; ensure you still set _externalWallet =
Web3.Wallet, create sessionAccount, compute SessionTokenPDA via
FindSessionToken, and populate SessionWallet fields only after these checks.
LocalAssociationScenario does not implement IDisposable in the current codebase. using var causes a compile error against main. Changed to plain var with TODO comments on lines 63 and 196. Will be updated to using var after PR magicblock-labs#260 merges, which adds IDisposable to LocalAssociationScenario. Related to magicblock-labs#267
There was a problem hiding this comment.
Actionable comments posted: 4
♻️ Duplicate comments (3)
Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs (3)
232-235:⚠️ Potential issue | 🟠 Major
OnWsErrormust unblock the waiters, otherwise an error-without-close hangs the scenario.
ConnectWithBackoffAsyncawaits_wsConnected.Task; the action loop awaits_responseTcs.Task. IfOnErrorfires without a subsequentOnClose(common for handshake/TLS failures), neither TCS is ever completed and the entire flow sits until the key-exchange / overall timeout expires. Making the handler instance-scoped and signaling the appropriate TCS (mirroringOnWsClose) fixes it.🛠️ Signal the appropriate TCS on error
- private static void OnWsError(string message) - { - Debug.Log($"[MWA] WS Error: {message}"); - } + private void OnWsError(string message) + { + Debug.Log($"[MWA] WS Error: {message}"); + if (_isConnecting) + _wsConnected?.TrySetResult(false); + else + { + var exc = new Exception($"WS error: {message}"); + _responseTcs?.TrySetException(exc); + _tcs?.TrySetException(exc); + } + }Also applies to Lines 177-184 where the connect loop relies on
_wsConnectedbeing signaled.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines 232 - 235, OnWsError currently only logs the error and can leave ConnectWithBackoffAsync (which awaits _wsConnected.Task) and the action loop (which awaits _responseTcs.Task) blocked if OnWsClose never fires; change the handler to be instance-scoped so it has access to the instance TCS fields and, inside OnWsError, mirror the logic from OnWsClose by setting/try-setting the appropriate TaskCompletionSources (_wsConnected and/or _responseTcs) to a failed or canceled state so waiters are unblocked (use TaskCompletionSource.TrySetException or TrySetCanceled as appropriate); ensure ConnectWithBackoffAsync and the connect loop will observe the signaled TCS when an error occurs.
53-56:⚠️ Potential issue | 🟠 MajorOverall timeout isn't applied to connect loop or action awaits.
cts.CancelAfter(_overallTimeout)is created, but_cancellationToken = ct(Line 56) andct.ThrowIfCancellationRequested()(Line 97) keep the original token, so the 30s ceiling never fires againstConnectWithBackoffAsync(used via_cancellationTokenat Lines 157/193) or the per-action loop when the caller passesCancellationToken.None.✅ Use the linked token everywhere
- _cancellationToken = ct; + _cancellationToken = cts.Token; @@ - ct.ThrowIfCancellationRequested(); + _cancellationToken.ThrowIfCancellationRequested();Also applies to Line 97.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines 53 - 56, The linked CancellationTokenSource (cts) created with CancellationTokenSource.CreateLinkedTokenSource(ct) and CancelAfter(_overallTimeout) isn't being used; replace usages of the original token with the linked token by assigning _cancellationToken = cts.Token (instead of ct), and ensure any ct.ThrowIfCancellationRequested() calls use cts.Token.ThrowIfCancellationRequested() (or call ThrowIfCancellationRequested on _cancellationToken), and pass _cancellationToken into ConnectWithBackoffAsync and the per-action loop awaits so the overall timeout actually cancels the connect loop and action awaits.
143-143:⚠️ Potential issue | 🟠 MajorAvoid emitting raw association tokens / decrypted payloads to
Debug.Log.Line 143 logs the plaintext
associationToken, and Line 257 logs the decrypted JSON ($"[MWA] Received: {json}"). Both end up inlogcatand any attached crash reporters. Log only non-sensitive metadata (port, payload length, or a truncated hash).🔒 Redacted logging
- Debug.Log($"[MWA] Launched intent for port {port}, token {associationToken}"); + Debug.Log($"[MWA] Launched intent for port {port}"); @@ - Debug.Log($"[MWA] Received: {json}"); + Debug.Log($"[MWA] Received response (len={json.Length})");Also applies to Line 257.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` at line 143, The Debug.Log calls in LocalAssociationScenario are emitting sensitive plaintext (the associationToken variable and the json payload); update the log statements that reference associationToken and json so they do not print raw secrets — e.g., log only non-sensitive metadata such as port, payload length, or a truncated/hashed version of associationToken (or a fixed masked string) instead of the full token; locate the Debug.Log lines that use associationToken and json in the LocalAssociationScenario class and replace their messages accordingly.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs`:
- Around line 270-283: Replace the 200ms polling in WaitForKeyExchangeAsync with
a TaskCompletionSource-based signal: add a field _keyExchangeTcs
(TaskCompletionSource<bool>) that WaitForKeyExchangeAsync awaits with the
provided CancellationToken and _keyExchangeTimeout, and remove the Task.Run
loop; in OnWsMessage, when the branch that constructs _client completes
successfully call _keyExchangeTcs.TrySetResult(true) (and call TrySetException
on failures), ensuring you create/reset _keyExchangeTcs before expecting a key
exchange and handle already-completed/timeout cases so WaitForKeyExchangeAsync
throws TimeoutException when the timeout elapses using _keyExchangeTimeout and
honors ct.
- Around line 237-268: The OnWsMessage handler can throw before _responseTcs is
created, causing an NRE in the catch and masking the real error; change all
TrySetException/ TrySetResult calls to use null-conditional (e.g.
_responseTcs?.TrySetException(...), _responseTcs?.TrySetResult(...)) and, in the
pre-_client branch (inside the _client == null block where
GenerateSessionEcdhSecret and MobileWalletAdapterClient are invoked), on error
route the failure to the active waiter by calling _tcs?.TrySetException(ex) or
setting _wsConnected appropriately so the key-exchange waiter is notified; also
apply the same null-conditional guard for stray messages arriving between
actions when you call _responseTcs.TrySetResult/TrySetException.
- Around line 285-307: Dispose is currently fire-and-for-get because
IDisposable.Dispose calls CleanupAsync() without awaiting, causing races on
_webSocket and _disposed; change the class to support true async disposal by
implementing IAsyncDisposable and moving the cleanup logic into ValueTask
IAsyncDisposable.DisposeAsync that awaits CleanupAsync(), update
IDisposable.Dispose to either call
DisposeAsync().AsTask().GetAwaiter().GetResult() or set _disposed true and
synchronously block on CleanupAsync().GetAwaiter().GetResult() if you must keep
synchronous semantics; ensure CleanupAsync still awaits _webSocket.Close(),
flips _disposed before starting cleanup to make it idempotent, and remove the
fire-and-forget `_ = CleanupAsync()` call in IDisposable.Dispose so callers stop
returning before socket close completes.
- Around line 47-121: Change the actions delegate from synchronous
Action<IAdapterOperations> to an async-aware Func<IAdapterOperations, Task> so
exceptions aren't swallowed: update the StartAndExecute signature to accept
List<Func<IAdapterOperations, Task>>, change where each action is invoked in the
loop (currently action.Invoke(_client) in StartAndExecute) to await
action(_client), and update all call sites in SolanaMobileWalletAdapter (_Login,
_SignAllTransactions, SignMessage) to pass async lambdas directly (remove any
cast to Action). Ensure _responseTcs/_tcs behavior remains the same so thrown
exceptions propagate into the surrounding try/catch and fault the
TaskCompletionSources.
---
Duplicate comments:
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs`:
- Around line 232-235: OnWsError currently only logs the error and can leave
ConnectWithBackoffAsync (which awaits _wsConnected.Task) and the action loop
(which awaits _responseTcs.Task) blocked if OnWsClose never fires; change the
handler to be instance-scoped so it has access to the instance TCS fields and,
inside OnWsError, mirror the logic from OnWsClose by setting/try-setting the
appropriate TaskCompletionSources (_wsConnected and/or _responseTcs) to a failed
or canceled state so waiters are unblocked (use
TaskCompletionSource.TrySetException or TrySetCanceled as appropriate); ensure
ConnectWithBackoffAsync and the connect loop will observe the signaled TCS when
an error occurs.
- Around line 53-56: The linked CancellationTokenSource (cts) created with
CancellationTokenSource.CreateLinkedTokenSource(ct) and
CancelAfter(_overallTimeout) isn't being used; replace usages of the original
token with the linked token by assigning _cancellationToken = cts.Token (instead
of ct), and ensure any ct.ThrowIfCancellationRequested() calls use
cts.Token.ThrowIfCancellationRequested() (or call ThrowIfCancellationRequested
on _cancellationToken), and pass _cancellationToken into ConnectWithBackoffAsync
and the per-action loop awaits so the overall timeout actually cancels the
connect loop and action awaits.
- Line 143: The Debug.Log calls in LocalAssociationScenario are emitting
sensitive plaintext (the associationToken variable and the json payload); update
the log statements that reference associationToken and json so they do not print
raw secrets — e.g., log only non-sensitive metadata such as port, payload
length, or a truncated/hashed version of associationToken (or a fixed masked
string) instead of the full token; locate the Debug.Log lines that use
associationToken and json in the LocalAssociationScenario class and replace
their messages accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 5c3db80f-ec3f-4343-88bf-9a7499656d3b
📒 Files selected for processing (1)
Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs
| public async Task<Response<object>> StartAndExecute(List<Action<IAdapterOperations>> actions, | ||
| CancellationToken ct = default) | ||
| { | ||
| if (actions == null || actions.Count == 0) | ||
| throw new ArgumentException("Actions must be non-null and non-empty"); | ||
| _actions = new Queue<Action<IAdapterOperations>>(actions); | ||
| var intent = LocalAssociationIntentCreator.CreateAssociationIntent( | ||
| _session.AssociationToken, | ||
| _port); | ||
| throw new ArgumentException("Actions required"); | ||
|
|
||
| using var cts = CancellationTokenSource.CreateLinkedTokenSource(ct); | ||
| cts.CancelAfter(_overallTimeout); | ||
|
|
||
| _cancellationToken = ct; | ||
| _tcs = new TaskCompletionSource<Response<object>>(); | ||
|
|
||
| StartActivityForAssociation(_session.AssociationToken, _port); | ||
|
|
||
| Debug.Log("[MWA] Waiting for websocket connection"); | ||
| await Task.Run(async () => | ||
| { | ||
| try | ||
| { | ||
| Debug.Log("[MWA Connect Thread] Started"); | ||
| _isConnecting = true; | ||
|
|
||
| await ConnectWithBackoffAsync(); | ||
|
|
||
| Debug.Log("[MWA Connect Thread] Completed"); | ||
| _isConnecting = false; | ||
|
|
||
| var helloReq = _session.CreateHelloReq(); | ||
| await _webSocket.Send(helloReq); | ||
|
|
||
| Debug.Log("[MWA] Hello sent. Waiting for pubkey..."); | ||
|
|
||
| await WaitForKeyExchangeAsync(cts.Token); | ||
|
|
||
| Debug.Log("[MWA] Pubkey received, session is encrypted"); | ||
|
|
||
| var queue = new Queue<Action<IAdapterOperations>>(actions); | ||
| Response<object> lastResponse = null; | ||
|
|
||
| while (queue.Count > 0) | ||
| { | ||
| _responseTcs = new TaskCompletionSource<Response<object>>(); | ||
|
|
||
| var action = queue.Dequeue(); | ||
| Debug.Log($"[MWA] Invoking action {action.Method.Name}"); | ||
| action.Invoke(_client); | ||
|
|
||
| lastResponse = await _responseTcs.Task; | ||
| _responseTcs = null; | ||
|
|
||
| ct.ThrowIfCancellationRequested(); | ||
| } | ||
|
|
||
| _tcs.TrySetResult(lastResponse ?? new Response<object>()); | ||
| } | ||
| catch (OperationCanceledException) | ||
| { | ||
| _tcs.TrySetResult(new Response<object> | ||
| { | ||
| Error = new Response<object>.ResponseError { Message = "Timeout or cancelled" } | ||
| }); | ||
| } | ||
| catch (Exception ex) | ||
| { | ||
| Debug.Log($"[MWA] Association failed: {ex}"); | ||
| _tcs.TrySetException(ex); | ||
| } | ||
| finally | ||
| { | ||
| await CleanupAsync(); | ||
| } | ||
| }, cts.Token); | ||
|
|
||
| return await _tcs.Task; | ||
| } |
There was a problem hiding this comment.
🛠️ Refactor suggestion | 🟠 Major
List<Action<IAdapterOperations>> still enables async-void exception swallowing — the core rewrite goal is only half-done.
The PR description calls out "cleaner error handling", but the delegate remains Action<IAdapterOperations> and is invoked fire-and-forget at Line 92 (action.Invoke(_client);). All three call sites in SolanaMobileWalletAdapter pass async client => { … } lambdas, which bind to Action as async void: any exception thrown inside (e.g., a JSON-RPC failure, a Deauthorize error, a cancellation) escapes this scenario entirely instead of faulting _responseTcs / _tcs. The loop then blocks on _responseTcs.Task until the overall timeout, masking the real error.
Migrating to Func<IAdapterOperations, Task> and awaiting each action lets exceptions propagate into the surrounding try/catch and turn into _tcs.TrySetException, which is what the new cancellation/timeout plumbing is meant to rely on.
♻️ Suggested signature change
- public async Task<Response<object>> StartAndExecute(List<Action<IAdapterOperations>> actions,
- CancellationToken ct = default)
+ public async Task<Response<object>> StartAndExecute(
+ IReadOnlyList<Func<IAdapterOperations, Task>> actions,
+ CancellationToken ct = default)
@@
- var action = queue.Dequeue();
- Debug.Log($"[MWA] Invoking action {action.Method.Name}");
- action.Invoke(_client);
-
- lastResponse = await _responseTcs.Task;
+ var action = queue.Dequeue();
+ Debug.Log($"[MWA] Invoking action {action.Method.Name}");
+ await action(_client).ConfigureAwait(false);
+
+ lastResponse = await _responseTcs.Task;Callers in SolanaMobileWalletAdapter._Login / _SignAllTransactions / SignMessage then drop the async client => … → Action cast and pass the lambda directly.
Based on learnings: prior PR #269 intentionally left the async client => { … } lambdas as Action<IAdapterOperations> (async void) because PR #260 was expected to migrate the delegate type to Func<IAdapterOperations, Task> as part of the ExecuteNextAction rewrite; leaving it as Action here keeps the exception-swallowing behavior the learning flagged.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
47 - 121, Change the actions delegate from synchronous
Action<IAdapterOperations> to an async-aware Func<IAdapterOperations, Task> so
exceptions aren't swallowed: update the StartAndExecute signature to accept
List<Func<IAdapterOperations, Task>>, change where each action is invoked in the
loop (currently action.Invoke(_client) in StartAndExecute) to await
action(_client), and update all call sites in SolanaMobileWalletAdapter (_Login,
_SignAllTransactions, SignMessage) to pass async lambdas directly (remove any
cast to Action). Ensure _responseTcs/_tcs behavior remains the same so thrown
exceptions propagate into the surrounding try/catch and fault the
TaskCompletionSources.
| private void OnWsMessage(byte[] bytes) | ||
| { | ||
| try | ||
| { | ||
| _session.GenerateSessionEcdhSecret(m); | ||
| var messageSender = new MobileWalletAdapterWebSocket(_webSocket, _session); | ||
| _client = new MobileWalletAdapterClient(messageSender); | ||
| _webSocket.OnMessage -= ReceivePublicKeyHandler; | ||
| _webSocket.OnMessage += HandleEncryptedSessionPayload; | ||
|
|
||
| // Executing the first action | ||
| ExecuteNextAction(); | ||
| // First message expected: raw pubkey for ECDH | ||
| if (_client == null) | ||
| { | ||
| _session.GenerateSessionEcdhSecret(bytes); | ||
| var messageSender = new MobileWalletAdapterWebSocket(_webSocket, _session); | ||
| _client = new MobileWalletAdapterClient(messageSender); | ||
|
|
||
| Debug.Log("[MWA] Key exchange complete → encrypted session ready"); | ||
| } | ||
| // All other should be encrypted messages | ||
| else | ||
| { | ||
| var decrypted = _session.DecryptSessionPayload(bytes); | ||
| var json = System.Text.Encoding.UTF8.GetString(decrypted); | ||
| _client.Receive(json); | ||
|
|
||
| Debug.Log($"[MWA] Received: {json}"); | ||
|
|
||
| var response = JsonConvert.DeserializeObject<Response<object>>(json); | ||
| _responseTcs.TrySetResult(response); | ||
| } | ||
| } | ||
| catch (Exception e) | ||
| catch (Exception ex) | ||
| { | ||
| Console.WriteLine(e); | ||
| Debug.Log($"[MWA] Message handler error: {ex}"); | ||
| _responseTcs.TrySetException(ex); | ||
| } | ||
| } |
There was a problem hiding this comment.
_responseTcs is null during key exchange — exception path NREs and silently hangs.
At the point OnWsMessage first fires (_client == null, Lines 243-249), _responseTcs has not been created yet — the action loop only allocates it at Line 88 after WaitForKeyExchangeAsync returns. If GenerateSessionEcdhSecret / the client construction throws, the catch at Line 266 calls _responseTcs.TrySetException(ex) and raises a NullReferenceException from inside the websocket event callback. The real failure is swallowed, WaitForKeyExchangeAsync keeps polling until its 20s timeout, and the caller sees a generic timeout instead of the underlying crypto/parse error.
Line 223 already uses _responseTcs?.TrySet… — the handlers below should be symmetric, and the pre-_client branch should fail fast via _tcs/_wsConnected instead.
🐛 Route failures to the active waiter
- catch (Exception ex)
- {
- Debug.Log($"[MWA] Message handler error: {ex}");
- _responseTcs.TrySetException(ex);
- }
+ catch (Exception ex)
+ {
+ Debug.Log($"[MWA] Message handler error: {ex}");
+ if (_client == null)
+ {
+ // Key-exchange phase: fail the connect waiter and the outer scenario.
+ _wsConnected?.TrySetException(ex);
+ _tcs?.TrySetException(ex);
+ }
+ else
+ {
+ _responseTcs?.TrySetException(ex);
+ }
+ }Consider the same ?. guard on Line 260 for stray messages arriving between actions.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
237 - 268, The OnWsMessage handler can throw before _responseTcs is created,
causing an NRE in the catch and masking the real error; change all
TrySetException/ TrySetResult calls to use null-conditional (e.g.
_responseTcs?.TrySetException(...), _responseTcs?.TrySetResult(...)) and, in the
pre-_client branch (inside the _client == null block where
GenerateSessionEcdhSecret and MobileWalletAdapterClient are invoked), on error
route the failure to the active waiter by calling _tcs?.TrySetException(ex) or
setting _wsConnected appropriately so the key-exchange waiter is notified; also
apply the same null-conditional guard for stray messages arriving between
actions when you call _responseTcs.TrySetResult/TrySetException.
| private Task WaitForKeyExchangeAsync(CancellationToken ct) | ||
| { | ||
| if (_actions.Count == 0 || response is { Failed: true }) | ||
| CloseAssociation(response); | ||
| var action = _actions.Dequeue(); | ||
| action.Invoke(_client); | ||
| return Task.Run(async () => | ||
| { | ||
| var start = DateTime.UtcNow; | ||
| while (_client == null) | ||
| { | ||
| if (ct.IsCancellationRequested || DateTime.UtcNow - start > _keyExchangeTimeout) | ||
| throw new TimeoutException("Key exchange timed out"); | ||
|
|
||
| await Task.Delay(200, ct); | ||
| } | ||
| }, ct); | ||
| } |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Prefer a TCS signal over 200 ms polling for key exchange.
WaitForKeyExchangeAsync busy-polls _client != null every 200 ms; OnWsMessage already knows the exact instant the pubkey is processed and could complete a dedicated TaskCompletionSource<bool> instead. That also gives you a natural place to TrySetException on key-exchange failure (related to the NRE issue raised above), and avoids the 0–200 ms latency added per handshake.
♻️ Signal-based wait
- private Task WaitForKeyExchangeAsync(CancellationToken ct)
- {
- return Task.Run(async () =>
- {
- var start = DateTime.UtcNow;
- while (_client == null)
- {
- if (ct.IsCancellationRequested || DateTime.UtcNow - start > _keyExchangeTimeout)
- throw new TimeoutException("Key exchange timed out");
-
- await Task.Delay(200, ct);
- }
- }, ct);
- }
+ private readonly TaskCompletionSource<bool> _keyExchangeTcs = new();
+
+ private async Task WaitForKeyExchangeAsync(CancellationToken ct)
+ {
+ using var timeoutCts = CancellationTokenSource.CreateLinkedTokenSource(ct);
+ timeoutCts.CancelAfter(_keyExchangeTimeout);
+ using (timeoutCts.Token.Register(() => _keyExchangeTcs.TrySetException(
+ new TimeoutException("Key exchange timed out"))))
+ {
+ await _keyExchangeTcs.Task;
+ }
+ }And in OnWsMessage's _client == null branch, call _keyExchangeTcs.TrySetResult(true); after the client is constructed (or TrySetException in the catch).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
270 - 283, Replace the 200ms polling in WaitForKeyExchangeAsync with a
TaskCompletionSource-based signal: add a field _keyExchangeTcs
(TaskCompletionSource<bool>) that WaitForKeyExchangeAsync awaits with the
provided CancellationToken and _keyExchangeTimeout, and remove the Task.Run
loop; in OnWsMessage, when the branch that constructs _client completes
successfully call _keyExchangeTcs.TrySetResult(true) (and call TrySetException
on failures), ensuring you create/reset _keyExchangeTcs before expecting a key
exchange and handle already-completed/timeout cases so WaitForKeyExchangeAsync
throws TimeoutException when the timeout elapses using _keyExchangeTimeout and
honors ct.
| private async Task CleanupAsync() | ||
| { | ||
| if (_webSocket is { State: WebSocketState.Open }) | ||
| await _webSocket.Close(); | ||
|
|
||
| if (_webSocket != null) | ||
| { | ||
| _webSocket.OnOpen -= OnWsOpen; | ||
| _webSocket.OnMessage -= OnWsMessage; | ||
| _webSocket.OnError -= OnWsError; | ||
| _webSocket.OnClose -= OnWsClose; | ||
| _webSocket = null; | ||
| } | ||
|
|
||
| _client = null; | ||
| _disposed = true; | ||
| } | ||
|
|
||
| private async void CloseAssociation(Response<object> response) | ||
| void IDisposable.Dispose() | ||
| { | ||
| _webSocket.OnMessage -= HandleEncryptedSessionPayload; | ||
| _handledEncryptedMessage = true; | ||
| await _webSocket.Close(); | ||
| _startAssociationTaskCompletionSource.SetResult(response); | ||
| if (_disposed) return; | ||
| _ = CleanupAsync(); | ||
| } |
There was a problem hiding this comment.
Dispose() is fire-and-forget — using returns before the socket is actually closed, and any cleanup exception is swallowed.
_ = CleanupAsync(); at Line 306 completes Dispose synchronously while _webSocket.Close() is still in flight. For the using var localAssociationScenario = new LocalAssociationScenario() call sites in SolanaMobileWalletAdapter, that means the next _Login / SignMessage can create a fresh scenario while the previous socket is still draining on the wallet side (the exact class of bug this PR is trying to fix). In addition, because _disposed is only flipped at the end of CleanupAsync (Line 300), two rapid Dispose calls will both start concurrent cleanups that race on _webSocket = null.
Two viable fixes:
- Implement
IAsyncDisposableand switch call sites toawait using, which is the idiomatic pattern when cleanup is genuinely async. - Make
Dispose()synchronous by blocking on cleanup (e.g.,CleanupAsync().GetAwaiter().GetResult()), guarded by setting_disposed = trueup-front to make it idempotent.
♻️ Option 1: IAsyncDisposable
-public class LocalAssociationScenario : IDisposable
+public class LocalAssociationScenario : IAsyncDisposable
@@
- void IDisposable.Dispose()
- {
- if (_disposed) return;
- _ = CleanupAsync();
- }
+ public async ValueTask DisposeAsync()
+ {
+ if (_disposed) return;
+ _disposed = true;
+ await CleanupAsync();
+ }Then at each call site: await using var localAssociationScenario = new LocalAssociationScenario();.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@Runtime/codebase/SolanaMobileStack/LocalAssociationScenario.cs` around lines
285 - 307, Dispose is currently fire-and-for-get because IDisposable.Dispose
calls CleanupAsync() without awaiting, causing races on _webSocket and
_disposed; change the class to support true async disposal by implementing
IAsyncDisposable and moving the cleanup logic into ValueTask
IAsyncDisposable.DisposeAsync that awaits CleanupAsync(), update
IDisposable.Dispose to either call
DisposeAsync().AsTask().GetAwaiter().GetResult() or set _disposed true and
synchronously block on CleanupAsync().GetAwaiter().GetResult() if you must keep
synchronous semantics; ensure CleanupAsync still awaits _webSocket.Close(),
flips _disposed before starting cleanup to make it idempotent, and remove the
fire-and-forget `_ = CleanupAsync()` call in IDisposable.Dispose so callers stop
returning before socket close completes.
1 participant
Problem
Issues connecting to and disconnecting from wallet apps on android
Solution
Rewrote LocalAssociationScenario.cs
Summary by CodeRabbit
New Features
Improvements