fix(strands): give the model the activity's cancellation signal

typescript/strands-agents/README.md

@@ -95,7 +95,7 @@ Before its first model call, each `TemporalAgent` reads the model's `modelId` an
 
 A worker model that keeps the conversation server-side, such as `OpenAIModel` with `stateful: true`, gets the agent's `modelState`, such as the last response id, on each model call through `TemporalAgent`, and the agent clears its messages after each turn, as natively. As in Strands, such an agent cannot use a `conversationManager` or `contextManager`; because `TemporalAgent` learns that the model is stateful when it reads the config, it throws at the agent's first invocation instead of at construction. Catch that error, or the workflow task fails and is retried until you remove the manager and redeploy the worker. With no manager given, such an agent keeps Strands' default sliding window manager, which Strands picks before the agent reads the config, where natively it would use `NullConversationManager`. Through the plugin that manager has nothing to trim: the agent clears its messages after each turn, before the manager runs, and a context window overflow reaches the agent as a model error. A custom model that reports `stateful` but does not use `modelState` ran statelessly through earlier versions, with the agent keeping its messages; it now gets only each turn's new message, and an agent that gives it a manager now throws. Report `stateful` as false for such a model to keep the earlier behavior.
 
-`OpenAIModel` calls through `TemporalAgent`, stateful or not, currently work only with `@strands-agents/sdk` releases before 1.15.0. From 1.15.0, Strands passes an `AbortSignal` in the model's stream options, the signal reaches the model activity as a plain object, and the OpenAI client fails with `signal.addEventListener is not a function` before it sends a request. Temporal retries that activity until its retry policy or timeout runs out.
+`OpenAIModel` calls through `TemporalAgent`, stateful or not, work with current `@strands-agents/sdk` releases because the model activities give the model their own cancellation signal (see [Cancellation](#cancellation)). Earlier versions of this plugin passed the agent's signal, which reached the model activity as a plain object, so from `@strands-agents/sdk` 1.15.0 the OpenAI client failed with `signal.addEventListener is not a function` before it sent a request. During a rollout, such calls on a worker without this change fail and are retried until an upgraded worker takes them.
 
 Stateful model calls run as their own activities, `strandsInvokeStatefulModel` and `strandsInvokeStatefulModelStreaming`, chosen from the stateful flag that `strandsGetModelConfig` returns, so a run that uses a stateful model needs the workflow and the config activity from the same version. An earlier workflow that reads this version's config schedules `invokeModel`, and this version then fails to replay the run with a nondeterminism error. This version's workflow that reads an earlier config, which has no flag, keeps the stateless path for the rest of the run, and an upgraded worker fails each of its model calls with a non-retryable `StatefulModelOnStatelessPath` error. Changing a worker model between stateful and stateless, such as adding `stateful: true` to a deployed `OpenAIModel`, also applies only to new runs: an open run keeps the mode its config read recorded, and its model calls fail with `StatefulModelOnStatelessPath`, or with `StatelessModelOnStatefulPath` when the model became stateless. None of these runs recovers. The agent's `invoke()` throws an `ActivityFailure` whose cause has that type, so unless the workflow catches it, the run fails; reset a run you want to continue. Roll out with Worker Versioning, keeping each run's workflow and activity workers on one version, including any task queue set in `activityOptions`. Without it, stop the runs that use a stateful model, replace every worker on those task queues, and then start new runs.
 
@@ -128,6 +128,24 @@ Workflow code that catches these errors around `invoke()` by class sees the `App
 
 If your worker lists one of these errors in `workflowFailureErrorTypes`, the run used to fail with a retryable failure, so a workflow retry policy started another run. The plugin's failure is non-retryable, so that run is no longer retried.
 
+## Cancellation
+
+The model gets the model activity's cancellation signal as `cancelSignal`. Once the worker learns that the activity was cancelled, the signal is aborted. Only a model that passes `cancelSignal` to its client stops the call; the model activity then ends as cancelled. In `@strands-agents/sdk` 1.19.0, `OpenAIModel` and `BedrockModel` do this, but `AnthropicModel`, `GoogleModel` and `VercelModel` ignore `cancelSignal`. With a model that ignores it, the provider call keeps running after cancellation, and the activity can still complete with the model's result. The worker learns of the cancellation through the activity's heartbeats, which the model activities send only when `activityOptions.heartbeatTimeout` is set. Without a `heartbeatTimeout`, the model call runs to completion after its activity is cancelled. `agent.cancel()` in workflow code does not cancel the model activity.
+
+Through `TemporalAgent`, cancelling the workflow during a model call cancels the model activity. With the default `TRY_CANCEL` cancellation type, the workflow does not wait for the activity: the run is cancelled at once (see [Retries](#retries)), and the model call is aborted once the worker learns of the cancellation. With `WAIT_CANCELLATION_COMPLETED`, the workflow waits for the model activity to end. If it ends as cancelled, so does the run. If it completes instead, as it can with a model that ignores `cancelSignal` or without a `heartbeatTimeout`, the agent gets the model's result and the turn carries on, so the run can complete normally rather than as cancelled:
+
+```ts
+import { ActivityCancellationType } from '@temporalio/workflow';
+
+new TemporalAgent({
+  activityOptions: {
+    startToCloseTimeout: '60 seconds',
+    heartbeatTimeout: '10 seconds',
+    cancellationType: ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+  },
+});
+```
+
 ## Snapshots
 
 `TemporalAgent.takeSnapshot()` and `TemporalAgent.loadSnapshot()` throw. Temporal's event history already persists workflow state durably at a finer granularity than Strands snapshots, so calling either inside a workflow is redundant.

typescript/strands-agents/src/__tests__/test-strands.ts

@@ -23,7 +23,7 @@ import {
   TestWorkflowEnvironment,
 } from '@temporalio/test-helpers';
 import { bundleWorkflowCode, DefaultLogger, NativeConnection, Worker, type WorkerOptions } from '@temporalio/worker';
-import { DeterminismViolationError } from '@temporalio/workflow';
+import { ActivityCancellationType, DeterminismViolationError } from '@temporalio/workflow';
 import { WorkflowStreamClient } from '@temporalio/workflow-streams/client';
 import { STRANDS_MALFORMED_INTERRUPT_TYPE, StrandsPlugin } from '..';
 import { ModelActivity } from '../model-activity';
@@ -43,6 +43,8 @@ import {
   activityToolAgent,
   failActivities,
   approveSignal,
+  cancellableAgent,
+  cancellableModelCall,
   compressionAgent,
   defaultModelAgent,
   doneSignal,
@@ -649,12 +651,14 @@ test('a history whose code depends on when a handler finishes still replays', as
   );
 });
 
-test('invokeModel reconstructs Message instances from wire data', async (t) => {
+test('invokeModel gives the model Message instances and an AbortSignal', async (t) => {
   // Strands content-block and Message classes have lossy `toJSON()` methods.
   // After Temporal's default payload converter serializes activity input,
   // the receiving side sees plain Bedrock-shape objects without the `type`
   // discriminator. ModelActivity must call `Message.fromMessageData` so the
-  // user's `Model.stream` receives real Strands types.
+  // user's `Model.stream` receives real Strands types. The agent's
+  // `cancelSignal` arrives as a plain object too, and model clients such as
+  // OpenAI's call its `addEventListener`.
   t.timeout(60_000);
   const { createWorker, executeWorkflow } = helpers(t);
 
@@ -663,8 +667,10 @@ test('invokeModel reconstructs Message instances from wire data', async (t) => {
   // reachable from the assertion below.
   let receivedFirstMessage: unknown;
   let firstBlockIsTextBlock = false;
+  let receivedSignal: unknown;
   class AssertingStub extends StubModel {
     override async *stream(messages: Message[], options?: StreamOptions): AsyncIterable<ModelStreamEvent> {
+      receivedSignal = options?.cancelSignal;
       if (receivedFirstMessage === undefined) {
         receivedFirstMessage = messages[0];
         const first = messages[0];
@@ -688,8 +694,137 @@ test('invokeModel reconstructs Message instances from wire data', async (t) => {
 
   t.true(receivedFirstMessage instanceof Message, 'expected first message to be a Message instance');
   t.true(firstBlockIsTextBlock, 'expected first content block to be a TextBlock instance');
+  t.true(receivedSignal instanceof AbortSignal, 'expected cancelSignal to be an AbortSignal');
 });
 
+for (const streaming of [false, true]) {
+  const path = streaming ? 'invokeModelStreaming' : 'invokeModel';
+  test(`cancelling the workflow aborts the model call in ${path} and cancels the workflow`, async (t) => {
+    t.timeout(60_000);
+    const { createWorker, startWorkflow } = helpers(t);
+    let receivedSignal: unknown;
+    let markStarted!: () => void;
+    const started = new Promise<void>((resolve) => (markStarted = resolve));
+    // Waits for its cancel signal, then fails as OpenAI's client does: an error named `Error`.
+    class AbortingStub extends StubModel {
+      override async *stream(messages: Message[], options?: StreamOptions): AsyncIterable<ModelStreamEvent> {
+        receivedSignal = options?.cancelSignal;
+        const aborted = await new Promise<boolean>((resolve) => {
+          options?.cancelSignal?.addEventListener('abort', () => resolve(true));
+          markStarted();
+        });
+        if (aborted) throw new Error('Request was aborted.');
+        yield* super.stream(messages, options);
+      }
+    }
+
+    const worker = await createWorker({
+      plugins: [new StrandsPlugin({ models: { test: () => new AbortingStub([]) } })],
+    });
+    await worker.runUntil(async () => {
+      const handle = await startWorkflow(cancellableModelCall, { args: [streaming] });
+      await started;
+      await handle.cancel();
+      const err = await t.throwsAsync(handle.result(), { instanceOf: WorkflowFailedError });
+      t.true(err!.cause instanceof CancelledFailure, `expected a CancelledFailure, got ${err!.cause}`);
+      const { events } = await handle.fetchHistory();
+      t.true(events!.some((e) => e.activityTaskCanceledEventAttributes), 'expected ActivityTaskCanceled');
+      t.false(events!.some((e) => e.activityTaskFailedEventAttributes), 'expected no ActivityTaskFailed');
+      t.true(events!.some((e) => e.workflowExecutionCanceledEventAttributes), 'expected WorkflowExecutionCanceled');
+    });
+    t.true(receivedSignal instanceof AbortSignal && receivedSignal.aborted, 'expected the model signal to abort');
+  });
+}
+
+for (const cancellationType of [
+  ActivityCancellationType.TRY_CANCEL,
+  ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+]) {
+  test(`cancelling a TemporalAgent workflow with ${cancellationType} aborts the model call and cancels the workflow`, async (t) => {
+    t.timeout(60_000);
+    const { createWorker, startWorkflow } = helpers(t);
+    let markStarted!: () => void;
+    const started = new Promise<void>((resolve) => (markStarted = resolve));
+    let markAborted!: () => void;
+    const aborted = new Promise<void>((resolve) => (markAborted = resolve));
+    // Waits for its cancel signal, then fails as OpenAI's client does: an error named `Error`.
+    class AbortingStub extends StubModel {
+      // eslint-disable-next-line require-yield
+      override async *stream(_messages: Message[], options?: StreamOptions): AsyncIterable<ModelStreamEvent> {
+        await new Promise<void>((resolve) => {
+          options?.cancelSignal?.addEventListener('abort', () => resolve());
+          markStarted();
+        });
+        markAborted();
+        throw new Error('Request was aborted.');
+      }
+    }
+
+    const worker = await createWorker({
+      plugins: [new StrandsPlugin({ models: { test: () => new AbortingStub([]) } })],
+    });
+    await worker.runUntil(async () => {
+      const handle = await startWorkflow(cancellableAgent, { args: [cancellationType] });
+      await started;
+      await handle.cancel();
+      const err = await t.throwsAsync(handle.result(), { instanceOf: WorkflowFailedError });
+      t.true(err!.cause instanceof CancelledFailure, `expected a CancelledFailure, got ${err!.cause}`);
+      // The model's signal aborts while the worker runs, not because it shuts down.
+      await aborted;
+      const { events } = await handle.fetchHistory();
+      t.true(events!.some((e) => e.activityTaskCancelRequestedEventAttributes), 'expected ActivityTaskCancelRequested');
+      t.false(events!.some((e) => e.workflowTaskFailedEventAttributes), 'expected no WorkflowTaskFailed');
+      t.true(events!.some((e) => e.workflowExecutionCanceledEventAttributes), 'expected WorkflowExecutionCanceled');
+      if (cancellationType === ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) {
+        t.true(events!.some((e) => e.activityTaskCanceledEventAttributes), 'expected ActivityTaskCanceled');
+      }
+    });
+  });
+
+  test(`cancelling a TemporalAgent workflow with ${cancellationType} while the model ignores its signal`, async (t) => {
+    t.timeout(60_000);
+    const { createWorker, startWorkflow } = helpers(t);
+    let markStarted!: () => void;
+    const started = new Promise<void>((resolve) => (markStarted = resolve));
+    let markAborted!: () => void;
+    const aborted = new Promise<void>((resolve) => (markAborted = resolve));
+    // Keeps going once its cancel signal aborts, as models that ignore `cancelSignal` do.
+    class IgnoringStub extends StubModel {
+      override async *stream(messages: Message[], options?: StreamOptions): AsyncIterable<ModelStreamEvent> {
+        await new Promise<void>((resolve) => {
+          options?.cancelSignal?.addEventListener('abort', () => resolve());
+          markStarted();
+        });
+        markAborted();
+        yield* super.stream(messages, options);
+      }
+    }
+
+    const worker = await createWorker({
+      plugins: [new StrandsPlugin({ models: { test: () => new IgnoringStub([textTurn('late result')]) } })],
+    });
+    await worker.runUntil(async () => {
+      const handle = await startWorkflow(cancellableAgent, { args: [cancellationType] });
+      await started;
+      await handle.cancel();
+      if (cancellationType === ActivityCancellationType.WAIT_CANCELLATION_COMPLETED) {
+        // The workflow waits for the activity, which completes, so the run completes with its result.
+        t.is((await handle.result()).trim(), 'late result');
+        const { events } = await handle.fetchHistory();
+        t.true(events!.some((e) => e.activityTaskCompletedEventAttributes), 'expected ActivityTaskCompleted');
+        t.true(events!.some((e) => e.workflowExecutionCompletedEventAttributes), 'expected WorkflowExecutionCompleted');
+        t.false(events!.some((e) => e.workflowExecutionCanceledEventAttributes), 'expected no WorkflowExecutionCanceled');
+      } else {
+        const err = await t.throwsAsync(handle.result(), { instanceOf: WorkflowFailedError });
+        t.true(err!.cause instanceof CancelledFailure, `expected a CancelledFailure, got ${err!.cause}`);
+        const { events } = await handle.fetchHistory();
+        t.true(events!.some((e) => e.workflowExecutionCanceledEventAttributes), 'expected WorkflowExecutionCanceled');
+      }
+      await aborted;
+    });
+  });
+}
+
 test('activityAsTool dispatches a registered activity from the agent loop', async (t) => {
   const { createWorker, executeWorkflow } = helpers(t);
 

typescript/strands-agents/src/__tests__/workflows/strands.ts

@@ -12,6 +12,7 @@ import {
   type Message,
 } from '@strands-agents/sdk';
 import {
+  ActivityCancellationType,
   ActivityFailure,
   allHandlersFinished,
   ApplicationFailure,
@@ -25,14 +26,34 @@ import {
   type ActivityOptions,
 } from '@temporalio/workflow';
 import { WorkflowStream } from '@temporalio/workflow-streams/workflow';
-import { TemporalAgent, TemporalMCPClient, type TemporalAgentOptions, workflow as strandsWorkflow } from '../..';
+import {
+  TemporalAgent,
+  TemporalMCPClient,
+  TemporalModel,
+  type TemporalAgentOptions,
+  workflow as strandsWorkflow,
+} from '../..';
 
 export async function helloAgent(prompt: string): Promise<string> {
   const agent = new TemporalAgent({ model: 'test', printer: false });
   const result = await agent.invoke(prompt);
   return result.toString();
 }
 
+/**
+ * Runs one agent turn whose model activity heartbeats, with the given cancellation type, so a
+ * workflow cancel during the model call reaches the model.
+ */
+export async function cancellableAgent(cancellationType: ActivityCancellationType): Promise<string> {
+  const agent = new TemporalAgent({
+    model: 'test',
+    printer: false,
+    activityOptions: { startToCloseTimeout: '30 seconds', heartbeatTimeout: '5 seconds', cancellationType },
+  });
+  const result = await agent.invoke('hello');
+  return result.toString();
+}
+
 /**
  * Runs each prompt as a turn of one agent with Strands' proactive compression, then returns the
  * config the agent's model reports in workflow code.
@@ -400,3 +421,25 @@ export async function streamingAgent(prompt: string): Promise<string> {
   const result = await agent.invoke(prompt);
   return result.toString();
 }
+
+/**
+ * Streams one model call whose activity heartbeats and waits for its cancellation to complete, so
+ * a workflow cancel reaches the model call and the workflow ends with the activity's result.
+ */
+export async function cancellableModelCall(streaming: boolean): Promise<string[]> {
+  if (streaming) void new WorkflowStream();
+  const model = new TemporalModel({
+    modelName: 'test',
+    streamingTopic: streaming ? 'model' : undefined,
+    activityOptions: {
+      startToCloseTimeout: '30 seconds',
+      heartbeatTimeout: '5 seconds',
+      cancellationType: ActivityCancellationType.WAIT_CANCELLATION_COMPLETED,
+    },
+  });
+  const types: string[] = [];
+  for await (const event of model.stream([])) {
+    types.push(event.type);
+  }
+  return types;
+}