Skip to content

[Brook] Added a Workflow Task interceptor that runs around each completion. - #198

Closed
moetemp wants to merge 2 commits into
moe/brook/v3/c0-upstream-mainfrom
moe/brook/v3/c1-workflow-task-interceptor
Closed

moetemp wants to merge 2 commits into
moe/brook/v3/c0-upstream-mainfrom
moe/brook/v3/c1-workflow-task-interceptor

Conversation

@moetemp

@moetemp moetemp commented Oct 21, 2026

Copy link
Copy Markdown
Owner

What changed?

This PR adds a public WorkflowTaskInterceptor trait and a WorkerExtension field on WorkerConfig that carries it.

Core awaits before_complete(info) before every RespondWorkflowTaskCompleted it sends for a run. That includes the local activity heartbeat completion and the completion of output buffered before an eviction, which have no lang activation. The hook site is in handle_activation_completed_success, before complete_workflow_task, which is the only production call of that RPC. An error, or a hook that misses info.deadline, fails the workflow task and evicts the run. Core reports the failure only for a first attempt, as it does today.

after_complete(info, accepted) runs after the answer. accepted is false on an RPC error and when the answer has reset_history_event_id > 0. Core waits at most 5 s on it. on_task_failed runs when Core fails a task, and on_evicted runs when the run leaves the cache, after after_complete.

The info has the run, started_event_id, attempt, started_time (poll field 13, now kept on ValidPollWFTQResponse), closes_run and a deadline at 90% of the workflow task timeout, measured from poll receipt. The hooks of a run block that run's next workflow task. The heartbeat completion runs inside the activation poll, so while its hooks run, that poll gives no activation to other runs. The replay worker clears the extension, so it never calls the hooks.

Why?

Core has no hook before it sends a workflow task completion. So a plugin cannot write data for a task before the server sees it, or act once the server accepts it. The trait is generic and Core knows nothing about what the plugin writes.

How did you test it?

Link to a test plan if any -

  • Unit Tests
  • Staging
  • End to End Tests

Mock worker tests in core_tests/workflow_task_interceptor.rs log the hooks and each completion the mock server gets, and check their order on each path: a normal completion, a hook error, a missed deadline, an UnhandledCommand refusal, a dropped speculative task, the local activity heartbeat, an eviction with output, a hung after_complete, and the replay worker.

A plugin can now run code before Core sends a workflow task completion and after the server answers. The hooks also run for the completions with no lang activation.
A hung after_complete could stall its run, and a task that waited on a busy run got a deadline past the server timeout.
@moetemp moetemp changed the title [Brook] Added a workflow task interceptor that runs around each completion. [Brook] Added a Workflow Task interceptor that runs around each completion. Oct 21, 2026
@moetemp

moetemp commented Oct 25, 2026

Copy link
Copy Markdown
Owner Author

Closing as stale. This layer is folded into the v5 tip, #209, and the upstream review uses one fresh PR instead. Plan: https://www.notion.so/4048fc567738814990b7c796ac28fe3f

@moetemp moetemp closed this Oct 25, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant