Skip to content
Open
Show file tree
Hide file tree
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
4 changes: 4 additions & 0 deletions apps/execution-worker/.env.example
Original file line number Diff line number Diff line change
Expand Up @@ -9,3 +9,7 @@ AI_MODEL=mistralai/mistral-small-3.2-24b-instruct
# free key at https://tavily.com (free tier ~1000 searches/month). Leave empty
# to disable: agents with web search toggled on still run, just without the tool.
TAVILY_API_KEY=

# Queue polled only by the specialized worker (start:specialized). A node type is
# routed here via its activity profile's `taskQueue`; unset means the default queue.
SPECIALIZED_TASK_QUEUE=workflow-execution-specialized
15 changes: 15 additions & 0 deletions apps/execution-worker/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -44,13 +44,28 @@ src/
└── engines/
└── temporal/
├── worker.ts # Worker bootstrap: executors + store, handed to WorkflowBuilderPlugin
├── specialized-worker.ts # Activity-only worker for a taskQueue-routed subset of node types
└── workflows.ts # One-line re-export of runWorkflow for Temporal's bundler
```

The workflow itself, the activity contract and the event emitter live in
[`@workflowbuilder/temporal`](../../packages/temporal/README.md). This app only supplies what is its
own: one executor per node type and the database as the store port.

## Per-node-type task queue routing

A node type's activity profile can carry `taskQueue`, which pins it to a queue other
than the default (`plugin.taskQueue`). `worker.ts` keeps polling the default queue for
everything else; `specialized-worker.ts` is a second, activity-only entrypoint (no
`workflowsPath` — Temporal supports activity-only workers) that polls
`SPECIALIZED_TASK_QUEUE` and registers only the node type(s) routed there. Run it with
`pnpm --filter execution-worker start:specialized`, or as the `worker-specialized`
compose service (see `deploy/ai-studio/README.md`).

This is a deployment-only change: `runGraph` and the graph model never see a taskQueue,
they only affect which worker process a node's `executeNode` activity is scheduled on.
A profile with no `taskQueue` behaves exactly as before.

## Temporal specifics

- **Task queue:** `workflow-execution`, read from `plugin.taskQueue` so the backend and the worker cannot drift apart. Both default to the same constant in the package.
Expand Down
1 change: 1 addition & 0 deletions apps/execution-worker/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@
"dev": "tsx watch --env-file=.env ./src/engines/temporal/worker.ts",
"start": "tsx --env-file=.env ./src/engines/temporal/worker.ts",
"start:prod": "tsx ./src/engines/temporal/worker.ts",
"start:specialized": "tsx ./src/engines/temporal/specialized-worker.ts",
"typecheck": "tsc --noEmit",
"lint": "eslint",
"lint:fix": "eslint --fix",
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,44 @@
// Activity-only worker: polls SPECIALIZED_TASK_QUEUE and registers just the node
// type(s) routed there via a profile's `taskQueue`. No `workflowsPath` — Temporal
// supports activity-only workers, so this process never needs the workflow bundle.
// Runs from its own Docker image so its tools stay off the general worker's image.
import { NativeConnection, Worker } from '@temporalio/worker';
import { WorkflowBuilderPlugin } from '@workflowbuilder/temporal';
import 'dotenv/config';

import { executeAiAgent } from '../../activities/ai-agent';
import { database } from '../../database';
import type { AiAgentNode } from '../../domain/ai-studio-nodes';
import { env } from '../../env';
import { logger } from '../../logger';
import { withPayloadSizeWarning } from '../../store-payload-warning';

const { createOpenRouter } = await import('@openrouter/ai-sdk-provider');

const openrouter = createOpenRouter({ apiKey: env.OPENROUTER_API_KEY });
const model = openrouter.chat(env.AI_MODEL);

const aiAgentLogger = logger.child({ component: 'ai-agent', worker: 'specialized' });

// Illustrative: the AI agent is the reference example of a node type heavy enough to
// want its own image (a coding-agent CLI, GPU tooling, ...). Swap in whichever
// executor(s) a real deployment routes to this queue via nodeActivityProfiles.
const plugin = new WorkflowBuilderPlugin<AiAgentNode>({
executors: {
'ai-studio/ai-agent': (node, context) =>
executeAiAgent(node, context, { model, logger: aiAgentLogger, tavilyApiKey: env.TAVILY_API_KEY }),
},
store: withPayloadSizeWarning(database, logger),
taskQueue: env.SPECIALIZED_TASK_QUEUE,
});

const connection = await NativeConnection.connect({ address: env.TEMPORAL_ADDRESS });

const worker = await Worker.create({
connection,
taskQueue: plugin.taskQueue,
plugins: [plugin],
});

logger.info('specialized execution worker started', { taskQueue: plugin.taskQueue });
await worker.run();
3 changes: 3 additions & 0 deletions apps/execution-worker/src/env.ts
Original file line number Diff line number Diff line change
Expand Up @@ -21,4 +21,7 @@ export const env = {
AI_MODEL: envOr('AI_MODEL', 'mistralai/mistral-small-3.2-24b-instruct'),
// Optional. Enables the AI Agent's web-search tool; agents run without it when unset.
TAVILY_API_KEY: process.env['TAVILY_API_KEY'],
// Polled only by the specialized worker (see engines/temporal/specialized-worker.ts).
// A distinct queue name, not the plugin's own, so routing a node here is opt-in per type.
SPECIALIZED_TASK_QUEUE: envOr('SPECIALIZED_TASK_QUEUE', 'workflow-execution-specialized'),
};
41 changes: 28 additions & 13 deletions deploy/ai-studio/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,21 +5,36 @@ any Docker host — an Azure VM, AWS, on-prem — with no cloud-specific glue.

## What runs

| Service | Image | Role | Exposed |
| ------------- | ------------------------------ | ----------------------------------------------- | ------------------------ |
| `web` | `ai-studio-web` (nginx) | Serves the SPA, proxies `/api` to the backend | `${WEB_PORT}` (only one) |
| `backend` | `ai-studio-runtime` | Hono REST + SSE event stream | internal |
| `worker` | `ai-studio-runtime` | Temporal worker, makes the OpenRouter LLM calls | internal |
| `temporal` | `temporalio/auto-setup` pinned | Workflow engine | internal |
| `app-db` | `postgres:16` | Workflow snapshots + execution events | internal |
| `temporal-db` | `postgres:16` | Temporal's own state store | internal |
| `temporal-ui` | `temporalio/ui` pinned | Debug only (`--profile debug`) | `127.0.0.1:8233` |
| Service | Image | Role | Exposed |
| -------------------- | ------------------------------ | ---------------------------------------------------------------------- | ------------------------ |
| `web` | `ai-studio-web` (nginx) | Serves the SPA, proxies `/api` to the backend | `${WEB_PORT}` (only one) |
| `backend` | `ai-studio-runtime` | Hono REST + SSE event stream | internal |
| `worker` | `ai-studio-runtime` | Temporal worker, makes the OpenRouter LLM calls | internal |
| `worker-specialized` | `ai-studio-runtime` | Activity-only worker for node types routed to `SPECIALIZED_TASK_QUEUE` | internal |
| `temporal` | `temporalio/auto-setup` pinned | Workflow engine | internal |
| `app-db` | `postgres:16` | Workflow snapshots + execution events | internal |
| `temporal-db` | `postgres:16` | Temporal's own state store | internal |
| `temporal-ui` | `temporalio/ui` pinned | Debug only (`--profile debug`) | `127.0.0.1:8233` |

Both images build from one Dockerfile (`deploy/ai-studio/Dockerfile`) with the
repo root as context. Backend and worker share a single image and differ only
in the compose `command`. Database migrations are applied by the backend at
boot (drizzle-orm's programmatic migrator) — there is no separate migration
service or step.
repo root as context. Backend, worker and worker-specialized share a single
image and differ only in the compose `command`. Database migrations are
applied by the backend at boot (drizzle-orm's programmatic migrator) — there
is no separate migration service or step.

`worker-specialized` polls its own task queue and registers only the node
type(s) meant to run on it (see `apps/execution-worker/README.md` § "Per-node-type
task queue routing"). It only executes anything once a node's activity profile
sets `taskQueue: 'workflow-execution-specialized'` (or your own
`SPECIALIZED_TASK_QUEUE`); until then it idles, polling a queue nothing is scheduled to.

### Running the specialized worker standalone

```bash
pnpm --filter execution-worker start:specialized # needs infra up + SPECIALIZED_TASK_QUEUE set
```

Or in compose alone: `docker compose up -d worker-specialized`.

## Quick start

Expand Down
24 changes: 24 additions & 0 deletions deploy/ai-studio/docker-compose.yml
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,30 @@ services:
condition: service_started
restart: unless-stopped

# dedicated pool for node types routed via a profile's `taskQueue`, so their
# dependencies (heavier tools, a different runtime) never reach the shared image.
# Same runtime image today — give it its own Dockerfile stage only once a routed
# node type actually needs extra OS-level tooling the shared image doesn't have.
worker-specialized:
image: ai-studio-runtime
build: *runtime-build
command: ['pnpm', '--filter', 'execution-worker', 'start:specialized']
environment:
DATABASE_URL: postgresql://wb:${APP_DB_PASSWORD:-wb}@app-db:5432/workflow_builder
TEMPORAL_ADDRESS: temporal:7233
OPENROUTER_API_KEY: ${OPENROUTER_API_KEY:?set OPENROUTER_API_KEY in deploy/ai-studio/.env}
AI_MODEL: ${AI_MODEL:-mistralai/mistral-small-3.2-24b-instruct}
TAVILY_API_KEY: ${TAVILY_API_KEY:-}
SPECIALIZED_TASK_QUEUE: ${SPECIALIZED_TASK_QUEUE:-workflow-execution-specialized}
depends_on:
app-db:
condition: service_healthy
backend:
condition: service_healthy
temporal:
condition: service_started
restart: unless-stopped

web:
image: ai-studio-web
build:
Expand Down
17 changes: 17 additions & 0 deletions packages/temporal/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -95,6 +95,23 @@ Keep the export named `runWorkflow`: that is the name the client starts, and a t

Entries are whole profiles rather than partials on purpose. A partial would let you set a timeout and silently drop the retry cap, and what Temporal falls back to is unlimited retries with backoff, which on a permanently failing model call is an unbounded bill. A node type with no entry resolves to `DEFAULT_NODE_ACTIVITY_PROFILE` and nothing else.

A profile can also carry `taskQueue`, routing that node type's activities to a worker other than the one running the workflow — a small pool with its own Docker image and tools, while every other node type keeps running on the default queue. Nothing else about the workflow changes: `runGraph` and the graph model never see it, it only decides which worker process picks up the `executeNode` command. Without a `taskQueue`, a profile behaves exactly as it did before this field existed.

```ts
export const nodeActivityProfiles: NodeActivityProfiles = {
// Needs a coding-agent CLI installed — its own image, its own worker.
'my-product/coding-agent': {
startToCloseTimeout: '30m',
retry: { maximumAttempts: 2 },
taskQueue: 'coding-agent-pool',
},
};
```

A worker for that queue is a second, activity-only `Worker.create` call (no `workflowsPath` — Temporal supports activity-only workers) registering just that node type's executor. See `apps/execution-worker/README.md` § "Per-node-type task queue routing" for a full example.

A `taskQueue` routed to a queue nobody polls fails silently at the workflow level: the activity is scheduled and simply never picked up. Checking that is only possible worker-side, against whichever queues a deployment knows it runs — `findProfilesWithUnpolledTaskQueue(profiles, polledTaskQueues)` from `/workflow` is the same shape as `findProfilesWithoutExecutor`, for the same reason: the sandbox has no registry to check it against itself.

A `startToCloseTimeout` is a number followed by `ms`, `s`, `m`, `h` or `d`. Decimals are fine (`'1.5h'`). It has to fit a protobuf `Duration`, so anything under one nanosecond or over `'3652500d'` is out. Zero, negative values and exponent notation are rejected even though TypeScript's template literal type admits them: `'0s'` type-checks, and Temporal treats a zero timeout as unset and refuses to schedule the activity.

This grammar is narrower than Temporal's own, which parses durations with the `ms` package and also takes `'30 minutes'` or `'1 week'`. One documented form is deliberate. If you think in the wider grammar, convert before the value reaches this map.
Expand Down
4 changes: 4 additions & 0 deletions packages/temporal/src/workflow/activity-profiles.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ export type DurationString = `${number}${'ms' | 's' | 'm' | 'h' | 'd'}`;
export type ActivityProfile = {
startToCloseTimeout: DurationString;
retry: { maximumAttempts: number };
// Routes this node type's activities to a non-default queue, e.g. a specialized
// worker with its own image/tools. Absent means the plugin's own task queue.
taskQueue?: string;
};

// Only for the two frozen singletons below. Annotating them `ActivityProfile` would
Expand All @@ -20,6 +23,7 @@ export type ActivityProfile = {
type ReadonlyActivityProfile = {
readonly startToCloseTimeout: DurationString;
readonly retry: { readonly maximumAttempts: number };
readonly taskQueue?: string;
};

// Node activities may call LLMs (minutes) — generous timeout, fewer retries to limit
Expand Down
34 changes: 34 additions & 0 deletions packages/temporal/src/workflow/node-activity-options.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -129,6 +129,40 @@ describe('resolveNodeActivityOptions', () => {
});
});

describe('taskQueue', () => {
it('forwards taskQueue when the profile carries one', () => {
const profiles: NodeActivityProfiles = {
'test/step': { startToCloseTimeout: '2m', retry: { maximumAttempts: 4 }, taskQueue: 'specialized' },
};

expect(resolveNodeActivityOptions(node(), profiles)).toEqual({
startToCloseTimeout: '2m',
retry: { maximumAttempts: 4 },
taskQueue: 'specialized',
});
});

it('omits the key entirely when the profile has none, not taskQueue: undefined', () => {
const resolved = resolveNodeActivityOptions(node(), {});

expect('taskQueue' in resolved).toBe(false);
});

it('is pure routing metadata: identical graph-relevant fields whether taskQueue is set or not', () => {
// The workflow sandbox must not branch on taskQueue — it only reaches proxyActivities.
const base = resolveNodeActivityOptions(node(), {
'test/step': { startToCloseTimeout: '2m', retry: { maximumAttempts: 4 } },
});
const routed = resolveNodeActivityOptions(node(), {
'test/step': { startToCloseTimeout: '2m', retry: { maximumAttempts: 4 }, taskQueue: 'specialized' },
});

expect(routed.startToCloseTimeout).toBe(base.startToCloseTimeout);
expect(routed.retry).toEqual(base.retry);
expect(routed.summary).toBe(base.summary);
});
});

describe('summary', () => {
it('carries the node label so Event History reads like the diagram', () => {
const resolved = resolveNodeActivityOptions(node({ label: 'Fetch order' }), {});
Expand Down
1 change: 1 addition & 0 deletions packages/temporal/src/workflow/node-activity-options.ts
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@ export function resolveFromValidatedProfiles(node: BaseNode, profiles: NodeActiv
const options: NodeActivityOptions = {
startToCloseTimeout: profile.startToCloseTimeout,
retry: { maximumAttempts: profile.retry.maximumAttempts },
...(profile.taskQueue === undefined ? {} : { taskQueue: profile.taskQueue }),
};

// Not trusted from the caller: any consumer can build the workflow input. Temporal
Expand Down
58 changes: 56 additions & 2 deletions packages/temporal/src/workflow/profile-validation.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import type { BaseNode } from './core-contract';
import { resolveNodeActivityOptions } from './node-activity-options';
import {
assertNodeActivityProfiles,
findProfilesWithUnpolledTaskQueue,
findProfilesWithoutExecutor,
freezeNodeActivityProfiles,
} from './profile-validation';
Expand Down Expand Up @@ -129,13 +130,43 @@ describe('assertNodeActivityProfiles', () => {
'test/step': {
startToCloseTimeout: '10m',
retry: { maximumAttempts: 2 },
taskQueue: 'x',
scheduleToCloseTimeout: '1h',
heartbeatTimeout: '1m',
},
};

expect(() => assertNodeActivityProfiles(extra as unknown as NodeActivityProfiles)).toThrow(
/has unknown keys "taskQueue", "heartbeatTimeout"/,
/has unknown keys "scheduleToCloseTimeout", "heartbeatTimeout"/,
);
});
});

describe('taskQueue', () => {
it('accepts a profile with a non-empty taskQueue', () => {
const withQueue = {
'test/step': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 }, taskQueue: 'specialized' },
};

expect(() => assertNodeActivityProfiles(withQueue as unknown as NodeActivityProfiles)).not.toThrow();
});

it('accepts a profile with no taskQueue, unchanged from before', () => {
expect(() => assertNodeActivityProfiles(profiles('10m'))).not.toThrow();
});

it('rejects an empty string, naming the path', () => {
const empty = { 'test/step': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 }, taskQueue: '' } };

expect(() => assertNodeActivityProfiles(empty as unknown as NodeActivityProfiles)).toThrow(
/nodeActivityProfiles\["test\/step"\]\.taskQueue must be a non-empty string/,
);
});

it('rejects a non-string value', () => {
const wrongType = { 'test/step': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 }, taskQueue: 7 } };

expect(() => assertNodeActivityProfiles(wrongType as unknown as NodeActivityProfiles)).toThrow(
/taskQueue must be a non-empty string/,
);
});
});
Expand Down Expand Up @@ -205,3 +236,26 @@ describe('findProfilesWithoutExecutor', () => {
expect(findProfilesWithoutExecutor(inherited, executors)).toEqual(['constructor']);
});
});

describe('findProfilesWithUnpolledTaskQueue', () => {
it('names a node type routed to a queue nobody polls', () => {
const map = {
'test/routed': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 }, taskQueue: 'specialized' },
'test/default': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 } },
} as unknown as NodeActivityProfiles;

expect(findProfilesWithUnpolledTaskQueue(map, new Set(['workflow-execution']))).toEqual(['test/routed']);
});

it('is silent once the queue is in the polled set', () => {
const map = {
'test/routed': { startToCloseTimeout: '10m', retry: { maximumAttempts: 2 }, taskQueue: 'specialized' },
} as unknown as NodeActivityProfiles;

expect(findProfilesWithUnpolledTaskQueue(map, new Set(['specialized']))).toEqual([]);
});

it('ignores a profile with no taskQueue regardless of the polled set', () => {
expect(findProfilesWithUnpolledTaskQueue(profiles('10m'), new Set())).toEqual([]);
});
});
Loading