Skip to content

Commit 4a137f6

Browse files
vincentkocvincentkoc
committed
fix(node): negotiate session attribution envelopes
1 parent 2c53be9 commit 4a137f6

2 files changed

Lines changed: 561 additions & 10 deletions

File tree

‎src/OpenClaw.Shared/WindowsNodeClient.cs‎

Lines changed: 225 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System;
2+
using System.Collections.Concurrent;
23
using System.Collections.Frozen;
34
using System.Collections.Generic;
45
using System.Diagnostics;
@@ -17,6 +18,11 @@ namespace OpenClaw.Shared;
1718
/// </summary>
1819
public class WindowsNodeClient : WebSocketClientBase
1920
{
21+
private const string NodeInvokeSessionKeyEnvelopeProtocolFeature =
22+
"node-invoke-session-key-envelope-v1";
23+
private static readonly TimeSpan s_protocolFeatureNegotiationTimeout =
24+
TimeSpan.FromSeconds(5);
25+
2026
private readonly DeviceIdentity _deviceIdentity;
2127

2228
// Node capabilities registry
@@ -60,6 +66,14 @@ public class WindowsNodeClient : WebSocketClientBase
6066
// at 8. When full, the gateway receives an immediate "node busy, retry" error response.
6167
private readonly SemaphoreSlim _invokeSemaphore = new(8, 8);
6268
private readonly InvocationCancellationRegistry _activeInvocations = new();
69+
private readonly ConcurrentDictionary<
70+
string,
71+
TaskCompletionSource<NodeInvokeSessionEnvelopeMode>> _protocolFeatureRequests = new();
72+
private Task<NodeInvokeSessionEnvelopeMode> _nodeInvokeSessionEnvelopeMode =
73+
Task.FromResult(NodeInvokeSessionEnvelopeMode.Authoritative);
74+
private int _gatewayConnectionGeneration;
75+
private readonly object _nodeInvokeEventDispatchLock = new();
76+
private Task _nodeInvokeEventDispatch = Task.CompletedTask;
6377

6478
// Events
6579
public event EventHandler<NodeInvokeRequest>? InvokeReceived;
@@ -318,10 +332,18 @@ private async Task HandleEventAsync(JsonElement root)
318332
await HandlePairingResolvedEventAsync(root, eventType);
319333
break;
320334
case "node.invoke.request":
321-
await HandleNodeInvokeEventAsync(root);
335+
QueueNodeInvokeMessage(
336+
root.Clone(),
337+
QueuedNodeInvokeKind.Invoke,
338+
containerName: "payload",
339+
responseId: null);
322340
break;
323341
case "node.invoke.cancel":
324-
await HandleNodeInvokeCancelAsync(root, "payload", responseId: null);
342+
QueueNodeInvokeMessage(
343+
root.Clone(),
344+
QueuedNodeInvokeKind.Cancel,
345+
containerName: "payload",
346+
responseId: null);
325347
break;
326348
case "health":
327349
if (root.TryGetProperty("payload", out var payload))
@@ -422,7 +444,67 @@ private async Task HandlePairingResolvedEventAsync(JsonElement root, string? eve
422444
}
423445
}
424446

425-
private async Task HandleNodeInvokeEventAsync(JsonElement root)
447+
private void QueueNodeInvokeMessage(
448+
JsonElement root,
449+
QueuedNodeInvokeKind kind,
450+
string containerName,
451+
string? responseId)
452+
{
453+
var envelopeModeTask = _nodeInvokeSessionEnvelopeMode;
454+
var connectionGeneration = _gatewayConnectionGeneration;
455+
lock (_nodeInvokeEventDispatchLock)
456+
{
457+
_nodeInvokeEventDispatch = _nodeInvokeEventDispatch
458+
.ContinueWith(
459+
_ => DispatchNodeInvokeEventAsync(
460+
root,
461+
kind,
462+
containerName,
463+
responseId,
464+
envelopeModeTask,
465+
connectionGeneration),
466+
CancellationToken.None,
467+
TaskContinuationOptions.ExecuteSynchronously,
468+
TaskScheduler.Default)
469+
.Unwrap();
470+
}
471+
}
472+
473+
private async Task DispatchNodeInvokeEventAsync(
474+
JsonElement root,
475+
QueuedNodeInvokeKind kind,
476+
string containerName,
477+
string? responseId,
478+
Task<NodeInvokeSessionEnvelopeMode> envelopeModeTask,
479+
int connectionGeneration)
480+
{
481+
try
482+
{
483+
var envelopeMode = await envelopeModeTask;
484+
if (connectionGeneration != _gatewayConnectionGeneration)
485+
return;
486+
487+
switch (kind)
488+
{
489+
case QueuedNodeInvokeKind.Invoke:
490+
// This returns after active-invocation registration and Task.Run handoff;
491+
// the next queued cancel is not serialized behind capability execution.
492+
await StartNodeInvokeEventAsync(root, envelopeMode);
493+
break;
494+
case QueuedNodeInvokeKind.Cancel:
495+
await HandleNodeInvokeCancelAsync(root, containerName, responseId);
496+
break;
497+
}
498+
}
499+
catch (Exception ex)
500+
{
501+
_logger.Error("Node invoke event dispatch failed", ex);
502+
}
503+
}
504+
505+
private async Task StartNodeInvokeEventAsync(
506+
JsonElement root,
507+
NodeInvokeSessionEnvelopeMode envelopeMode)
426508
{
427509
var telemetry = new NodeToolInvocation(NodeToolTransport.Gateway);
428510
_logger.Info("[NODE] Received node.invoke.request event");
@@ -508,7 +590,7 @@ await SendGatewayResultAndCompleteTelemetryAsync(
508590
}
509591
}
510592

511-
var sessionKey = ExtractGatewayNodeInvokeSessionKey(payload);
593+
var sessionKey = ExtractGatewayNodeInvokeSessionKey(payload, args, envelopeMode);
512594

513595
_logger.Info($"[NODE] Invoking command: {command}");
514596

@@ -731,6 +813,22 @@ private string BuildNodeConnectMessage(string? nonce, long? challengeTimestampMs
731813

732814
internal void HandleResponse(JsonElement root)
733815
{
816+
if (root.TryGetProperty("id", out var idProp) &&
817+
idProp.ValueKind == JsonValueKind.String &&
818+
idProp.GetString() is { } responseId &&
819+
_protocolFeatureRequests.TryRemove(responseId, out var negotiation))
820+
{
821+
var mode = IsUnsupportedNodeProtocolFeaturesResponse(root)
822+
? NodeInvokeSessionEnvelopeMode.Legacy
823+
: NodeInvokeSessionEnvelopeMode.Authoritative;
824+
negotiation.TrySetResult(mode);
825+
if (mode == NodeInvokeSessionEnvelopeMode.Legacy)
826+
{
827+
_logger.Debug("[NODE] Gateway does not accept node protocol feature publication");
828+
}
829+
return;
830+
}
831+
734832
if (root.TryGetProperty("ok", out var okProp) &&
735833
okProp.ValueKind == JsonValueKind.False)
736834
{
@@ -748,6 +846,7 @@ internal void HandleResponse(JsonElement root)
748846
if (payload.TryGetProperty("type", out var t) && t.GetString() == "hello-ok")
749847
{
750848
_logger.Info("[HANDSHAKE] Received hello-ok!");
849+
_gatewayConnectionGeneration++;
751850
PublishGatewaySelf(GatewaySelfInfo.FromHelloOk(payload));
752851
var reconnectingAfterApproval = _pairingApprovedAwaitingReconnect;
753852
_isConnected = true;
@@ -827,10 +926,75 @@ internal void HandleResponse(JsonElement root)
827926
}
828927

829928
RaiseStatusChanged(ConnectionStatus.Connected);
929+
_ = PublishNodeProtocolFeaturesAsync();
830930
HandshakeSucceeded?.Invoke(this, EventArgs.Empty);
831931
}
832932
}
833933

934+
private async Task PublishNodeProtocolFeaturesAsync()
935+
{
936+
var requestId = Guid.NewGuid().ToString();
937+
var negotiation = new TaskCompletionSource<NodeInvokeSessionEnvelopeMode>(
938+
TaskCreationOptions.RunContinuationsAsynchronously);
939+
_nodeInvokeSessionEnvelopeMode = negotiation.Task;
940+
var request = new
941+
{
942+
type = "req",
943+
id = requestId,
944+
method = "node.protocolFeatures.update",
945+
@params = new
946+
{
947+
features = new[] { NodeInvokeSessionKeyEnvelopeProtocolFeature }
948+
}
949+
};
950+
951+
try
952+
{
953+
_protocolFeatureRequests.TryAdd(requestId, negotiation);
954+
_ = ResolveNodeProtocolFeaturesTimeoutAsync(requestId, negotiation);
955+
await SendRawAsync(JsonSerializer.Serialize(request));
956+
}
957+
catch (Exception ex)
958+
{
959+
_protocolFeatureRequests.TryRemove(requestId, out _);
960+
negotiation.TrySetResult(NodeInvokeSessionEnvelopeMode.Authoritative);
961+
// Only an explicit unknown-method response enables legacy attribution.
962+
// Transport failures keep omitted envelopes fail-closed.
963+
_logger.Warn($"[NODE] Failed to publish protocol features: {ex.Message}");
964+
}
965+
}
966+
967+
private async Task ResolveNodeProtocolFeaturesTimeoutAsync(
968+
string requestId,
969+
TaskCompletionSource<NodeInvokeSessionEnvelopeMode> negotiation)
970+
{
971+
await Task.Delay(s_protocolFeatureNegotiationTimeout);
972+
if (_protocolFeatureRequests.TryGetValue(requestId, out var pending) &&
973+
ReferenceEquals(pending, negotiation))
974+
{
975+
_logger.Warn("[NODE] Protocol feature publication timed out; using fail-closed envelopes");
976+
negotiation.TrySetResult(NodeInvokeSessionEnvelopeMode.Authoritative);
977+
}
978+
}
979+
980+
private static bool IsUnsupportedNodeProtocolFeaturesResponse(JsonElement response)
981+
{
982+
if (!response.TryGetProperty("ok", out var ok) ||
983+
ok.ValueKind != JsonValueKind.False ||
984+
!response.TryGetProperty("error", out var error) ||
985+
error.ValueKind != JsonValueKind.Object)
986+
{
987+
return false;
988+
}
989+
990+
return error.TryGetProperty("code", out var code) &&
991+
string.Equals(code.GetString(), "INVALID_REQUEST", StringComparison.OrdinalIgnoreCase) &&
992+
error.TryGetProperty("message", out var message) &&
993+
message.GetString()?.Contains(
994+
"unknown method: node.protocolFeatures.update",
995+
StringComparison.Ordinal) == true;
996+
}
997+
834998
private void TryStoreHandshakeDeviceToken(string token, string[]? scopes)
835999
{
8361000
try
@@ -1119,7 +1283,11 @@ private async Task HandleRequestAsync(JsonElement root)
11191283
await HandleNodeInvokeAsync(root, id);
11201284
break;
11211285
case "node.invoke.cancel":
1122-
await HandleNodeInvokeCancelAsync(root, "params", id);
1286+
QueueNodeInvokeMessage(
1287+
root.Clone(),
1288+
QueuedNodeInvokeKind.Cancel,
1289+
containerName: "params",
1290+
responseId: id);
11231291
break;
11241292
case "ping":
11251293
await SendPongAsync(id);
@@ -1568,19 +1736,48 @@ private static bool TryGetCancellationTargetId(JsonElement container, out string
15681736
return false;
15691737
}
15701738

1571-
private static string? ExtractGatewayNodeInvokeSessionKey(JsonElement envelope)
1739+
private static string? ExtractGatewayNodeInvokeSessionKey(
1740+
JsonElement envelope,
1741+
JsonElement args,
1742+
NodeInvokeSessionEnvelopeMode envelopeMode)
15721743
{
1573-
if (envelope.TryGetProperty("sessionKey", out var envelopeSessionKey) &&
1574-
envelopeSessionKey.ValueKind == JsonValueKind.String)
1744+
if (envelope.TryGetProperty("sessionKey", out var envelopeSessionKey))
15751745
{
1576-
var sessionKey = envelopeSessionKey.GetString();
1746+
if (envelopeSessionKey.ValueKind == JsonValueKind.String)
1747+
{
1748+
var sessionKey = envelopeSessionKey.GetString();
1749+
if (!string.IsNullOrWhiteSpace(sessionKey))
1750+
return sessionKey;
1751+
}
1752+
1753+
return null;
1754+
}
1755+
1756+
if (envelopeMode == NodeInvokeSessionEnvelopeMode.Legacy &&
1757+
args.ValueKind == JsonValueKind.Object &&
1758+
args.TryGetProperty("sessionKey", out var legacySessionKey) &&
1759+
legacySessionKey.ValueKind == JsonValueKind.String)
1760+
{
1761+
var sessionKey = legacySessionKey.GetString();
15771762
if (!string.IsNullOrWhiteSpace(sessionKey))
15781763
return sessionKey;
15791764
}
15801765

15811766
return null;
15821767
}
15831768

1769+
private enum NodeInvokeSessionEnvelopeMode
1770+
{
1771+
Authoritative,
1772+
Legacy
1773+
}
1774+
1775+
private enum QueuedNodeInvokeKind
1776+
{
1777+
Invoke,
1778+
Cancel
1779+
}
1780+
15841781
private sealed record CommandDispatchEntry(
15851782
INodeCapability Capability,
15861783
string CanonicalName);
@@ -1732,6 +1929,7 @@ protected override bool ShouldAutoReconnect()
17321929
protected override void OnDisconnected()
17331930
{
17341931
_activeInvocations.CancelAll();
1932+
ResetNodeInvokeSessionEnvelopeNegotiation();
17351933
_isConnected = false;
17361934
// Don't reset pairing state when disconnected due to pairing — gateway
17371935
// closes the socket after PAIRING_REQUIRED but we're still waiting for approval
@@ -1745,6 +1943,7 @@ protected override void OnDisconnected()
17451943
protected override void OnError(Exception ex)
17461944
{
17471945
_activeInvocations.CancelAll();
1946+
ResetNodeInvokeSessionEnvelopeNegotiation();
17481947
_isConnected = false;
17491948
if (!_pairingBlocked)
17501949
{
@@ -1756,5 +1955,22 @@ protected override void OnError(Exception ex)
17561955
protected override void OnDisposing()
17571956
{
17581957
_activeInvocations.CancelAll();
1958+
ResetNodeInvokeSessionEnvelopeNegotiation();
1959+
}
1960+
1961+
private void ResetNodeInvokeSessionEnvelopeNegotiation()
1962+
{
1963+
_gatewayConnectionGeneration++;
1964+
foreach (var request in _protocolFeatureRequests.Values)
1965+
{
1966+
request.TrySetResult(NodeInvokeSessionEnvelopeMode.Authoritative);
1967+
}
1968+
_protocolFeatureRequests.Clear();
1969+
_nodeInvokeSessionEnvelopeMode =
1970+
Task.FromResult(NodeInvokeSessionEnvelopeMode.Authoritative);
1971+
lock (_nodeInvokeEventDispatchLock)
1972+
{
1973+
_nodeInvokeEventDispatch = Task.CompletedTask;
1974+
}
17591975
}
17601976
}

0 commit comments

Comments
 (0)