Skip to content
Draft
Changes from all commits
Commits
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
71 changes: 60 additions & 11 deletions cli/beamable.common/Runtime/Promise.cs
Original file line number Diff line number Diff line change
Expand Up @@ -450,7 +450,19 @@ void ICriticalNotifyCompletion.UnsafeOnCompleted(Action continuation)

void INotifyCompletion.OnCompleted(Action continuation)
{
((ICriticalNotifyCompletion)this).UnsafeOnCompleted(continuation);
// Mirror TaskAwaiter.OnCompleted: flow ExecutionContext so AsyncLocal
// values set by the caller survive the await and are restored on the
// thread that resolves the promise.
var capturedContext = System.Threading.ExecutionContext.Capture();
if (capturedContext == null)
{
((ICriticalNotifyCompletion)this).UnsafeOnCompleted(continuation);
return;
}
((ICriticalNotifyCompletion)this).UnsafeOnCompleted(() =>
{
System.Threading.ExecutionContext.Run(capturedContext, s => ((Action)s)(), continuation);
});
Comment on lines +453 to +465
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

System.Threading.ExecutionContext is referenced here, but this file defines DISABLE_THREADING for UNITY_WEBGL and elsewhere conditionally compiles out System.Threading usage. As-is, this new code path can break WebGL builds (or any build with DISABLE_THREADING) because ExecutionContext lives in System.Threading. Consider wrapping the ExecutionContext capture/run logic in #if !DISABLE_THREADING and falling back to the previous UnsafeOnCompleted(continuation) behavior when threading is disabled.

Copilot uses AI. Check for mistakes.
}

/// <summary>
Expand Down Expand Up @@ -1253,6 +1265,18 @@ public void SetStateMachine(IAsyncStateMachine machine)
_stateMachine = machine;
}

private void MoveNextWithCapturedContext(System.Threading.ExecutionContext capturedContext)
{
if (capturedContext == null)
{
_stateMachine.MoveNext();
}
else
{
System.Threading.ExecutionContext.Run(capturedContext, s => ((IAsyncStateMachine)s).MoveNext(), _stateMachine);
}
Comment on lines +1268 to +1277
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This helper uses System.Threading.ExecutionContext, which conflicts with the existing DISABLE_THREADING compilation mode (defined for UNITY_WEBGL at the top of this file) that is meant to avoid System.Threading APIs. Please gate ExecutionContext usage behind #if !DISABLE_THREADING (or equivalent) and keep a non-threading fallback that just calls _stateMachine.MoveNext() so builds with DISABLE_THREADING still compile.

Copilot uses AI. Check for mistakes.
}

public void AwaitOnCompleted<TAwaiter, TStateMachine>(
ref TAwaiter awaiter, ref TStateMachine stateMachine)
where TAwaiter : INotifyCompletion
Expand All @@ -1264,18 +1288,26 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(
_stateMachine.SetStateMachine(stateMachine);
}

awaiter.OnCompleted(() =>
{
_stateMachine.MoveNext();
});
// Mirror AsyncTaskMethodBuilder: capture ExecutionContext at the await
// point and restore it before resuming the state machine so AsyncLocal
// values flow across the await.
var capturedContext = System.Threading.ExecutionContext.Capture();
awaiter.OnCompleted(() => MoveNextWithCapturedContext(capturedContext));
Comment on lines +1291 to +1295
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This change alters core await/continuation behavior (ExecutionContext/AsyncLocal flow) but there doesn’t appear to be a unit test that asserts AsyncLocal (or logging scope) values survive await Promise boundaries. Given there are existing PromiseTests, please add a focused test that sets an AsyncLocal value, awaits a Promise that completes on a different callback, and verifies the value is present after resumption (and ideally that suppression/clearing behaves as expected).

Copilot uses AI. Check for mistakes.
}

public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
ref TAwaiter awaiter, ref TStateMachine stateMachine)
where TAwaiter : ICriticalNotifyCompletion
where TStateMachine : IAsyncStateMachine
{
AwaitOnCompleted(ref awaiter, ref stateMachine);
if (_stateMachine == null)
{
_stateMachine = stateMachine;
_stateMachine.SetStateMachine(stateMachine);
}

var capturedContext = System.Threading.ExecutionContext.Capture();
awaiter.UnsafeOnCompleted(() => MoveNextWithCapturedContext(capturedContext));
Comment on lines +1303 to +1310
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

System.Threading.ExecutionContext is captured/used in the async method builder here, but this file supports a DISABLE_THREADING build mode (notably for WebGL) where System.Threading APIs are avoided. Please add conditional compilation so the builder still compiles when DISABLE_THREADING is defined (eg, skip capture and resume the state machine directly).

Copilot uses AI. Check for mistakes.
}

public void Start<TStateMachine>(ref TStateMachine stateMachine)
Expand Down Expand Up @@ -1313,6 +1345,18 @@ public void SetStateMachine(IAsyncStateMachine machine)
_stateMachine = machine;
}

private void MoveNextWithCapturedContext(System.Threading.ExecutionContext capturedContext)
{
if (capturedContext == null)
{
_stateMachine.MoveNext();
}
else
{
System.Threading.ExecutionContext.Run(capturedContext, s => ((IAsyncStateMachine)s).MoveNext(), _stateMachine);
}
}

public void AwaitOnCompleted<TAwaiter, TStateMachine>(
ref TAwaiter awaiter, ref TStateMachine stateMachine)
where TAwaiter : INotifyCompletion
Expand All @@ -1324,18 +1368,23 @@ public void AwaitOnCompleted<TAwaiter, TStateMachine>(
_stateMachine.SetStateMachine(stateMachine);
}

awaiter.OnCompleted(() =>
{
_stateMachine.MoveNext();
});
var capturedContext = System.Threading.ExecutionContext.Capture();
awaiter.OnCompleted(() => MoveNextWithCapturedContext(capturedContext));
}

public void AwaitUnsafeOnCompleted<TAwaiter, TStateMachine>(
ref TAwaiter awaiter, ref TStateMachine stateMachine)
where TAwaiter : ICriticalNotifyCompletion
where TStateMachine : IAsyncStateMachine
{
AwaitOnCompleted(ref awaiter, ref stateMachine);
if (_stateMachine == null)
{
_stateMachine = stateMachine;
_stateMachine.SetStateMachine(stateMachine);
}

var capturedContext = System.Threading.ExecutionContext.Capture();
awaiter.UnsafeOnCompleted(() => MoveNextWithCapturedContext(capturedContext));
Comment on lines +1371 to +1387
Copy link

Copilot AI Apr 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Same ExecutionContext dependency exists in the non-generic async method builder; without conditional compilation this can break DISABLE_THREADING targets (eg UNITY_WEBGL). Please apply the same #if !DISABLE_THREADING guard/fallback here as well so this code remains portable across supported Unity platforms.

Copilot uses AI. Check for mistakes.
}

public void Start<TStateMachine>(ref TStateMachine stateMachine)
Expand Down
Loading