Skip to content

Pass single-session options to Matching task-add hooks #12840

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

Draft
Shivs11 wants to merge 4 commits into
base: shivam/agentcore-single-session
Choose a base branch
from
+640 −305
Draft

Pass single-session options to Matching task-add hooks #12840

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
f8dfae1
Pass single-session options to Matching task-add hooks
Shivs11 Oct 30, 2026
f753a99
Merge updated session state into WCI hooks
Shivs11 Nov 4, 2026
60ad752
Merge updated single-session base into WCI hooks
Shivs11 Nov 6, 2026
2d25964
Use an internal task session directive for History dispatch
Shivs11 Nov 6, 2026
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
525 changes: 275 additions & 250 deletions api/matchingservice/v1/request_response.pb.go
View file
Open in desktop

Large diffs are not rendered by default.

2 changes: 2 additions & 0 deletions client/matching/client_gen.go
View file
Open in desktop

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

5 changes: 5 additions & 0 deletions proto/internal/temporal/server/api/matchingservice/v1/request_response.proto
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@ import "temporal/api/history/v1/message.proto";
import "temporal/api/nexus/v1/message.proto";
import "temporal/api/protocol/v1/message.proto";
import "temporal/api/query/v1/message.proto";
import "temporal/api/session/v1/message.proto";
import "temporal/api/taskqueue/v1/message.proto";
import "temporal/api/worker/v1/message.proto";
import "temporal/api/workflowservice/v1/request_response.proto";
Expand Down Expand Up @@ -168,6 +169,8 @@ message AddWorkflowTaskRequest {
temporal.api.common.v1.Priority priority = 12;
// Stamp value from when the workflow task was scheduled. Used to validate the task is still relevant.
int32 stamp = 13;
// Identifies the bound single-session destination for this task, if any.
temporal.api.session.v1.SingleSessionOptions single_session_options = 14;

@ShahabT ShahabT Nov 3, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

it's kinda weird to me that the workflow (i.e. the session) deal with SingleSessionOptions and has to pass it as part of tasks to matching.

The mutable state and history->protocol should not use SingleSessionOptions IMO, because that proto is a worker options proto.

When we look from session/workflow POV they would have info/options about their "exclusive worker". When we look from the sticky TQ it has info about it's single assigned session. When looked from worker it has options to state single session preferences.

@Shivs11 Shivs11 Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

was thinking of the SingleSessionOptions being a proto that can be used in places where deemed necessary; i look at the worker options as the ones that are specifically defined on the worker, under WorkerSessionOptions.

@ShahabT ShahabT Nov 4, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

but this is the same SingleSessionOptions that is used in WorkerSessionOptions. I'm not sure the options wf wants to have for single-session mode will be always the same as the info that workflow wants to remember.

}

message AddWorkflowTaskResponse {
Expand Down Expand Up @@ -196,6 +199,8 @@ message AddActivityTaskRequest {
// necessary start information is carried within this component, obviating the need to use the fields that apply to
// embedded activities.
bytes component_ref = 14;
// Identifies the bound single-session destination for this task, if any.
temporal.api.session.v1.SingleSessionOptions single_session_options = 15;
Comment thread
Shivs11 marked this conversation as resolved.
Resolved
}

message AddActivityTaskResponse {
Expand Down
1 change: 1 addition & 0 deletions service/history/api/updateworkflow/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -423,6 +423,7 @@ func (u *Updater) addWorkflowTaskToMatching(ctx context.Context) error {
RunId: u.wfKey.RunID,
},
TaskQueue: u.taskQueue,
SingleSessionOptions: workflow.SingleSessionOptions(u.singleSessionBound, u.taskQueue),
ScheduledEventId: u.scheduledEventID,
ScheduleToStartTimeout: durationpb.New(u.scheduleToStartTimeout),
Clock: clock,
Expand Down
4 changes: 4 additions & 0 deletions service/history/ndc_standby_task_util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -184,6 +184,7 @@ type (
activityTaskScheduleToStartTimeout time.Duration
versionDirective *taskqueuespb.TaskVersionDirective
priority *commonpb.Priority
singleSessionBound bool
}

verifyCompletionRecordedPostActionInfo struct {
Expand All @@ -199,6 +200,7 @@ type (
taskqueue *taskqueuepb.TaskQueue
versionDirective *taskqueuespb.TaskVersionDirective
priority *commonpb.Priority
singleSessionBound bool
}
)

Expand Down Expand Up @@ -226,6 +228,7 @@ func newActivityTaskPostActionInfo(
activityTaskScheduleToStartTimeout: activityInfo.ScheduleToStartTimeout.AsDuration(),
versionDirective: directive,
priority: priority,
singleSessionBound: mutableState.GetExecutionInfo().GetSingleSessionBound(),
}, nil
}

Expand All @@ -242,6 +245,7 @@ func newWorkflowTaskPostActionInfo(
taskqueue: taskqueue,
versionDirective: directive,
priority: priority,
singleSessionBound: mutableState.GetExecutionInfo().GetSingleSessionBound(),
}, nil
}

Expand Down
2 changes: 2 additions & 0 deletions service/history/timer_queue_active_task_executor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -614,6 +614,7 @@ func (t *timerQueueActiveTaskExecutor) executeActivityRetryTimerTask(
// activity scheduling. Returning this error keeps retrying the timer task without dispatch.
return err
}
singleSessionBound := mutableState.GetExecutionInfo().GetSingleSessionBound()
scheduleToStartTimeout := timestamp.DurationValue(activityInfo.ScheduleToStartTimeout)
directive := MakeDirectiveForActivityTask(mutableState, activityInfo)
useWfBuildId := activityInfo.GetUseWorkflowBuildIdInfo() != nil
Expand All @@ -629,6 +630,7 @@ func (t *timerQueueActiveTaskExecutor) executeActivityRetryTimerTask(
RunId: task.GetRunID(),
},
TaskQueue: taskQueue,
SingleSessionOptions: workflow.SingleSessionOptions(singleSessionBound, taskQueue),
ScheduledEventId: task.EventID,
ScheduleToStartTimeout: durationpb.New(scheduleToStartTimeout),
Clock: vclock.NewVectorClock(t.shardContext.GetClusterMetadata().GetClusterID(), t.shardContext.GetShardID(), task.TaskID),
Expand Down
1 change: 1 addition & 0 deletions service/history/timer_queue_standby_task_executor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -788,6 +788,7 @@ func (t *timerQueueStandbyTaskExecutor) pushActivity(
RunId: activityTask.RunID,
},
TaskQueue: pushActivityInfo.taskQueue,
SingleSessionOptions: workflow.SingleSessionOptions(pushActivityInfo.singleSessionBound, pushActivityInfo.taskQueue),
ScheduledEventId: activityTask.EventID,
ScheduleToStartTimeout: durationpb.New(activityScheduleToStartTimeout),
Clock: vclock.NewVectorClock(t.shardContext.GetClusterMetadata().GetClusterID(), t.shardContext.GetShardID(), activityTask.TaskID),
Expand Down
5 changes: 4 additions & 1 deletion service/history/transfer_queue_active_task_executor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -278,6 +278,7 @@ func (t *transferQueueActiveTaskExecutor) processActivityTask(
if err != nil {
return err
}
singleSessionBound := mutableState.GetExecutionInfo().GetSingleSessionBound()
timeout := timestamp.DurationValue(ai.ScheduleToStartTimeout)
directive := MakeDirectiveForActivityTask(mutableState, ai)
priority := priorities.Merge(mutableState.GetExecutionInfo().Priority, ai.Priority)
Expand All @@ -287,7 +288,7 @@ func (t *transferQueueActiveTaskExecutor) processActivityTask(
// the rest of logic is making RPC call, which takes time.
release(nil)

return t.pushActivity(ctx, task, taskQueue, timeout, directive, priority, historyi.TransactionPolicyActive)
return t.pushActivity(ctx, task, taskQueue, singleSessionBound, timeout, directive, priority, historyi.TransactionPolicyActive)
}

func (t *transferQueueActiveTaskExecutor) processWorkflowTask(
Expand Down Expand Up @@ -345,6 +346,7 @@ func (t *transferQueueActiveTaskExecutor) processWorkflowTask(
ctx,
transferTask,
taskQueue,
singleSessionBound,
scheduleToStartTimeout.AsDuration(),
directive,
priority,
Expand All @@ -369,6 +371,7 @@ func (t *transferQueueActiveTaskExecutor) processWorkflowTask(
ctx,
transferTask,
taskQueue,
singleSessionBound,
scheduleToStartTimeout.AsDuration(),
directive,
priority,
Expand Down
10 changes: 10 additions & 0 deletions service/history/transfer_queue_active_task_executor_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ import (
"go.temporal.io/server/common/tasktoken"
"go.temporal.io/server/common/telemetry"
"go.temporal.io/server/common/testing/protomock"
"go.temporal.io/server/common/testing/protorequire"
"go.temporal.io/server/common/testing/testhooks"
"go.temporal.io/server/common/worker_versioning"
"go.temporal.io/server/service/history/configs"
Expand Down Expand Up @@ -254,9 +255,11 @@ func (s *transferQueueActiveTaskExecutorSuite) TestProcessActivityTask_Success()
name string
session bool
missingPlacement bool
bound bool
}{
{name: "regular activity with sticky workflow"},
{name: "workflow-session activity", session: true},
{name: "single-session activity", session: true, bound: true},
{name: "workflow-session activity awaiting placement", session: true, missingPlacement: true},
} {
s.Run(tc.name, func() {
Expand Down Expand Up @@ -302,6 +305,7 @@ func (s *transferQueueActiveTaskExecutorSuite) TestProcessActivityTask_Success()
if !tc.missingPlacement {
mutableState.SetStickyTaskQueue("sticky-worker", durationpb.New(time.Second))
}
mutableState.GetExecutionInfo().SingleSessionBound = tc.bound

transferTask := &tasks.ActivityTask{
WorkflowKey: definition.NewWorkflowKey(
Expand All @@ -324,6 +328,9 @@ func (s *transferQueueActiveTaskExecutorSuite) TestProcessActivityTask_Success()
request.TaskQueue = &taskqueuepb.TaskQueue{Name: "sticky-worker", Kind: enumspb.TASK_QUEUE_KIND_STICKY, NormalName: taskQueueName}
request.VersionDirective = nil
}
if tc.bound {
request.SingleSessionOptions = &sessionpb.SingleSessionOptions{StableIdentity: "sticky-worker"}
}
s.mockMatchingClient.EXPECT().AddActivityTask(gomock.Any(), protomock.Eq(request), gomock.Any()).Return(&matchingservice.AddActivityTaskResponse{}, nil)
}

Expand Down Expand Up @@ -719,6 +726,9 @@ func (s *transferQueueActiveTaskExecutorSuite) TestProcessWorkflowTask_Sticky_No
s.Equal(enumspb.TASK_QUEUE_KIND_STICKY, request.GetTaskQueue().GetKind())
if tc.bound {
s.Equal(5*time.Second, request.GetScheduleToStartTimeout().AsDuration())
if !tc.unavailable {
protorequire.ProtoEqual(s.T(), &sessionpb.SingleSessionOptions{StableIdentity: stickyTaskQueueName}, request.GetSingleSessionOptions())
}
}
if tc.unavailable {
return nil, &serviceerrors.StickyWorkerUnavailable{}
Expand Down
2 changes: 2 additions & 0 deletions service/history/transfer_queue_standby_task_executor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -711,6 +711,7 @@ func (t *transferQueueStandbyTaskExecutor) pushActivity(
ctx,
activityTask,
pushActivityInfo.taskQueue,
pushActivityInfo.singleSessionBound,
pushActivityInfo.activityTaskScheduleToStartTimeout,
pushActivityInfo.versionDirective,
pushActivityInfo.priority,
Expand All @@ -733,6 +734,7 @@ func (t *transferQueueStandbyTaskExecutor) pushWorkflowTask(
ctx,
task.(*tasks.WorkflowTask),
pushwtInfo.taskqueue,
pushwtInfo.singleSessionBound,
pushwtInfo.workflowTaskScheduleToStartTimeout,
pushwtInfo.versionDirective,
pushwtInfo.priority,
Expand Down
5 changes: 5 additions & 0 deletions service/history/transfer_queue_task_executor_base.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ import (
"go.temporal.io/server/service/history/queues"
"go.temporal.io/server/service/history/tasks"
"go.temporal.io/server/service/history/vclock"
"go.temporal.io/server/service/history/workflow"
wcache "go.temporal.io/server/service/history/workflow/cache"
"google.golang.org/protobuf/types/known/durationpb"
)
Expand Down Expand Up @@ -95,6 +96,7 @@ func (t *transferQueueTaskExecutorBase) pushActivity(
ctx context.Context,
task *tasks.ActivityTask,
taskQueue *taskqueuepb.TaskQueue,
singleSessionBound bool,
activityScheduleToStartTimeout time.Duration,
directive *taskqueuespb.TaskVersionDirective,
priority *commonpb.Priority,
Expand All @@ -107,6 +109,7 @@ func (t *transferQueueTaskExecutorBase) pushActivity(
RunId: task.RunID,
},
TaskQueue: taskQueue,
SingleSessionOptions: workflow.SingleSessionOptions(singleSessionBound, taskQueue),
ScheduledEventId: task.ScheduledEventID,
ScheduleToStartTimeout: durationpb.New(activityScheduleToStartTimeout),
Clock: vclock.NewVectorClock(t.shardContext.GetClusterMetadata().GetClusterID(), t.shardContext.GetShardID(), task.TaskID),
Expand Down Expand Up @@ -148,6 +151,7 @@ func (t *transferQueueTaskExecutorBase) pushWorkflowTask(
ctx context.Context,
task *tasks.WorkflowTask,
taskqueue *taskqueuepb.TaskQueue,
singleSessionBound bool,
workflowTaskScheduleToStartTimeout time.Duration,
directive *taskqueuespb.TaskVersionDirective,
priority *commonpb.Priority,
Expand All @@ -164,6 +168,7 @@ func (t *transferQueueTaskExecutorBase) pushWorkflowTask(
RunId: task.RunID,
},
TaskQueue: taskqueue,
SingleSessionOptions: workflow.SingleSessionOptions(singleSessionBound, taskqueue),
ScheduledEventId: task.ScheduledEventID,
ScheduleToStartTimeout: sst,
Clock: vclock.NewVectorClock(t.shardContext.GetClusterMetadata().GetClusterID(), t.shardContext.GetShardID(), task.TaskID),
Expand Down
9 changes: 9 additions & 0 deletions service/history/workflow/util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ import (
enumspb "go.temporal.io/api/enums/v1"
historypb "go.temporal.io/api/history/v1"
"go.temporal.io/api/serviceerror"
sessionpb "go.temporal.io/api/session/v1"
taskqueuepb "go.temporal.io/api/taskqueue/v1"
workflowpb "go.temporal.io/api/workflow/v1"
enumsspb "go.temporal.io/server/api/enums/v1"
persistencespb "go.temporal.io/server/api/persistence/v1"
Expand All @@ -26,6 +28,13 @@ import (
"google.golang.org/protobuf/types/known/timestamppb"
)

func SingleSessionOptions(singleSessionBound bool, taskQueue *taskqueuepb.TaskQueue) *sessionpb.SingleSessionOptions {
if !singleSessionBound || taskQueue.GetKind() != enumspb.TASK_QUEUE_KIND_STICKY || taskQueue.GetName() == "" {
return nil
}
return &sessionpb.SingleSessionOptions{StableIdentity: taskQueue.GetName()}
}

func failWorkflowTask(
mutableState historyi.MutableState,
workflowTask *historyi.WorkflowTaskInfo,
Expand Down
3 changes: 3 additions & 0 deletions service/matching/hooks/task_lifecycle_hooks.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (

deploymentpb "go.temporal.io/api/deployment/v1"
enumspb "go.temporal.io/api/enums/v1"
sessionpb "go.temporal.io/api/session/v1"
"go.temporal.io/server/common/namespace"
"go.temporal.io/server/common/tqid"
)
Expand Down Expand Up @@ -41,6 +42,8 @@ type (
DeploymentVersion *deploymentpb.WorkerDeploymentVersion
IsSyncMatch bool // Deprecated: use SyncMatchOutcome instead.
SyncMatchOutcome SyncMatchOutcome
// Non-nil only when the task belongs to an already-bound single session.
SingleSessionOptions *sessionpb.SingleSessionOptions
}

TaskHookFactory interface {
Expand Down
10 changes: 6 additions & 4 deletions service/matching/matching_engine.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -720,8 +720,9 @@ func (e *matchingEngineImpl) AddWorkflowTask(
}

return pm.AddTask(ctx, addTaskParams{
taskInfo: taskInfo,
forwardInfo: addRequest.ForwardInfo,
taskInfo: taskInfo,
forwardInfo: addRequest.ForwardInfo,
singleSessionOptions: addRequest.GetSingleSessionOptions(),
})
}

Expand Down Expand Up @@ -771,8 +772,9 @@ func (e *matchingEngineImpl) AddActivityTask(
}

return pm.AddTask(ctx, addTaskParams{
taskInfo: taskInfo,
forwardInfo: addRequest.ForwardInfo,
taskInfo: taskInfo,
forwardInfo: addRequest.ForwardInfo,
singleSessionOptions: addRequest.GetSingleSessionOptions(),
})
}

Expand Down
6 changes: 4 additions & 2 deletions service/matching/physical_task_queue_manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ import (

enumspb "go.temporal.io/api/enums/v1"
"go.temporal.io/api/serviceerror"
sessionpb "go.temporal.io/api/session/v1"
taskqueuepb "go.temporal.io/api/taskqueue/v1"
"go.temporal.io/api/workflowservice/v1"
enumsspb "go.temporal.io/server/api/enums/v1"
Expand Down Expand Up @@ -54,8 +55,9 @@ const (

type (
addTaskParams struct {
taskInfo *persistencespb.TaskInfo
forwardInfo *taskqueuespb.TaskForwardInfo
taskInfo *persistencespb.TaskInfo
forwardInfo *taskqueuespb.TaskForwardInfo
singleSessionOptions *sessionpb.SingleSessionOptions
}
// physicalTaskQueueManagerImpl manages a set of physical queues that comprise one logical
// queue, corresponding to a single versioned queue of a task queue partition.
Expand Down
Loading
Loading