@rivetkit/effect 0.0.0-06-09-refactor-rivetkit-split-workflow-context-into-workflowcontext-workflowstepcontext.e0c9540

package/src/State.test.ts

@@ -0,0 +1,181 @@
+import { assert, describe, it } from "@effect/vitest";
+import { State } from "@rivetkit/effect";
+import { Effect, Exit, PubSub, Stream } from "effect";
+
+// Helper: build a State backed by a plain mutable cell, with
+// Effect-typed read/write closures. Mirrors how Registry wires
+// `decodeUnknownEffect` / `encodeUnknownEffect` over `c.state`.
+const makeCellState = <A>(initial: A) => {
+	const cell = { value: initial };
+	return State.make<A, never, never>(
+		() => Effect.sync(() => cell.value),
+		(v) =>
+			Effect.sync(() => {
+				cell.value = v;
+			}),
+	).pipe(Effect.map((s) => ({ s, cell })));
+};
+
+describe("State", () => {
+	it.effect("get reflects the backing store", () =>
+		Effect.gen(function* () {
+			const { s, cell } = yield* makeCellState(42);
+			assert.strictEqual(yield* State.get(s), 42);
+
+			cell.value = 100;
+			assert.strictEqual(yield* State.get(s), 100);
+		}),
+	);
+
+	it.effect("set writes through to the backing store", () =>
+		Effect.gen(function* () {
+			const { s, cell } = yield* makeCellState(0);
+			yield* State.set(s, 7);
+			assert.strictEqual(cell.value, 7);
+			assert.strictEqual(yield* State.get(s), 7);
+		}),
+	);
+
+	it.effect("update applies f over read/write", () =>
+		Effect.gen(function* () {
+			const { s, cell } = yield* makeCellState(10);
+			yield* State.update(s, (n) => n + 5);
+			assert.strictEqual(cell.value, 15);
+		}),
+	);
+
+	it.effect("updateAndGet returns the new value and commits it", () =>
+		Effect.gen(function* () {
+			const { s, cell } = yield* makeCellState(10);
+			const next = yield* State.updateAndGet(s, (n) => n + 5);
+			assert.strictEqual(next, 15);
+			assert.strictEqual(cell.value, 15);
+		}),
+	);
+
+	it.effect("modify returns B and commits the new value", () =>
+		Effect.gen(function* () {
+			const { s, cell } = yield* makeCellState("a");
+			const b = yield* State.modify(
+				s,
+				(str) => [str.length, `${str}b`] as const,
+			);
+			assert.strictEqual(b, 1);
+			assert.strictEqual(cell.value, "ab");
+		}),
+	);
+
+	it.effect(
+		"update is atomic across concurrent fibers (no lost updates)",
+		() =>
+			Effect.gen(function* () {
+				const { s, cell } = yield* makeCellState(0);
+				yield* Effect.all(
+					Array.from({ length: 100 }, () =>
+						State.update(s, (n) => n + 1),
+					),
+					{ concurrency: "unbounded" },
+				);
+				assert.strictEqual(cell.value, 100);
+			}),
+	);
+
+	it.effect("changes replays the most recent published value", () =>
+		Effect.gen(function* () {
+			const { s } = yield* makeCellState(0);
+			const initial = yield* State.changes(s).pipe(
+				Stream.take(1),
+				Stream.runCollect,
+			);
+			assert.deepStrictEqual(initial, [0]);
+
+			State.publishUnsafe(s, 7);
+			const later = yield* State.changes(s).pipe(
+				Stream.take(1),
+				Stream.runCollect,
+			);
+			assert.deepStrictEqual(later, [7]);
+		}),
+	);
+
+	it.effect("publish pushes values to live subscribers", () =>
+		Effect.gen(function* () {
+			const { s } = yield* makeCellState(0);
+			yield* Effect.scoped(
+				Effect.gen(function* () {
+					const sub = yield* PubSub.subscribe(s.pubsub);
+					assert.strictEqual(yield* PubSub.take(sub), 0);
+
+					yield* State.publish(s, 1);
+					yield* State.publish(s, 2);
+					assert.strictEqual(yield* PubSub.take(sub), 1);
+					assert.strictEqual(yield* PubSub.take(sub), 2);
+				}),
+			);
+		}),
+	);
+
+	it.effect("set does NOT auto-publish — the runtime does", () =>
+		Effect.gen(function* () {
+			const { s } = yield* makeCellState(0);
+			yield* State.set(s, 99);
+			// replay should still hold the initial 0, not 99
+			const latest = yield* State.changes(s).pipe(
+				Stream.take(1),
+				Stream.runCollect,
+			);
+			assert.deepStrictEqual(latest, [0]);
+		}),
+	);
+
+	it.effect("isState discriminates", () =>
+		Effect.gen(function* () {
+			const { s } = yield* makeCellState(0);
+			assert.isTrue(State.isState(s));
+			assert.isFalse(State.isState({}));
+			assert.isFalse(State.isState(null));
+			assert.isFalse(State.isState(42));
+		}),
+	);
+
+	it.effect("supports .pipe()", () =>
+		Effect.gen(function* () {
+			const { s } = yield* makeCellState(0);
+			yield* s.pipe(State.set(5));
+			assert.strictEqual(yield* State.get(s), 5);
+
+			yield* s.pipe(State.update((n) => n * 2));
+			assert.strictEqual(yield* State.get(s), 10);
+		}),
+	);
+
+	it.effect("read failure propagates through get", () =>
+		Effect.gen(function* () {
+			const reads = { count: 0 };
+			// Construction reads once to seed the pubsub; subsequent reads
+			// fail. Mirrors a schema mismatch on persisted state.
+			const s = yield* State.make<number, "boom", never>(
+				() =>
+					Effect.suspend(() => {
+						reads.count++;
+						if (reads.count === 1) return Effect.succeed(0);
+						return Effect.fail("boom" as const);
+					}),
+				() => Effect.void,
+			);
+			const exit = yield* Effect.exit(State.get(s));
+			assert.isTrue(Exit.isFailure(exit));
+		}),
+	);
+
+	it.effect("write failure propagates through set", () =>
+		Effect.gen(function* () {
+			const s = yield* State.make<number, "boom", never>(
+				() => Effect.succeed(0),
+				() => Effect.fail("boom" as const),
+			);
+			const exit = yield* Effect.exit(State.set(s, 1));
+			assert.isTrue(Exit.isFailure(exit));
+		}),
+	);
+});

package/src/State.ts

@@ -0,0 +1,224 @@
+/**
+ * `State` is a typed view over an actor's persisted state, plus a
+ * subscribable stream of every change.
+ *
+ * Unlike a `Ref`, `State` has no in-memory cell — the persisted store
+ * is the source of truth. Reads decode the live store on demand;
+ * writes encode and overwrite it. A `PubSub<A>` backs {@link changes}
+ * and is fed externally — the runtime publishes to it from rivetkit's
+ * `onStateChange` callback so subscribers see every committed change,
+ * including ones initiated outside the SDK.
+ *
+ * Read and write are Effect-typed so schemas with asynchronous
+ * transforms (or service requirements) are supported. `update` and
+ * `modify` serialize through a per-`State` semaphore so read/apply/
+ * write triples are atomic across fibers; `set` shares the same lock
+ * so all writes are linearized.
+ *
+ * The PubSub uses replay = 1, matching `SubscriptionRef`: a new
+ * subscriber immediately sees the most recent value.
+ */
+import {
+	Effect,
+	Inspectable,
+	identity,
+	Pipeable,
+	Predicate,
+	PubSub,
+	Semaphore,
+	Stream,
+	type Types,
+} from "effect";
+import { dual } from "effect/Function";
+
+const TypeId = "~@rivetkit/effect/State";
+
+/**
+ * A view over a persisted state cell with a subscribable change stream.
+ *
+ * - `A` — the value type
+ * - `E` — the read/write closures' failure type (e.g. a schema's
+ *         `SchemaError` when read/write decode/encode against a schema)
+ * - `R` — the read/write closures' service requirements
+ */
+export interface State<A, E = never, R = never>
+	extends Variance<A, E, R>,
+		Pipeable.Pipeable,
+		Inspectable.Inspectable {
+	readonly read: () => Effect.Effect<A, E, R>;
+	readonly write: (value: A) => Effect.Effect<void, E, R>;
+	readonly pubsub: PubSub.PubSub<A>;
+	/**
+	 * Serializes writes (`set`, `update`, `modify`) so the read/apply/
+	 * write triple is atomic. The runtime may also use this semaphore
+	 * to serialize its own decode-and-publish work from
+	 * `onStateChange`, keeping the change stream's order consistent
+	 * with the write order.
+	 */
+	readonly semaphore: Semaphore.Semaphore;
+}
+
+export const isState = (u: unknown): u is State<unknown, unknown> =>
+	Predicate.hasProperty(u, TypeId);
+
+export interface Variance<A, E, R> {
+	readonly [TypeId]: {
+		readonly _A: Types.Invariant<A>;
+		readonly _E: Types.Covariant<E>;
+		readonly _R: Types.Covariant<R>;
+	};
+}
+
+const Proto = {
+	...Pipeable.Prototype,
+	...Inspectable.BaseProto,
+	[TypeId]: { _A: identity, _E: identity, _R: identity },
+	toJSON(this: State<unknown, unknown, unknown>) {
+		return { _id: "State" };
+	},
+};
+
+/**
+ * Creates a `State` from `read` and `write` closures over the
+ * underlying store. The closures are responsible for any
+ * encoding/decoding; `State` itself is schema-agnostic.
+ *
+ * The current value (per `read()`) is published to the pubsub on
+ * construction so any subscription obtained later replays it.
+ *
+ * The PubSub is not explicitly shut down — it's reclaimed by GC when
+ * the `State` and any subscribers become unreachable.
+ */
+export const make = Effect.fnUntraced(function* <A, E, R>(
+	read: () => Effect.Effect<A, E, R>,
+	write: (value: A) => Effect.Effect<void, E, R>,
+): Effect.fn.Return<State<A, E, R>, E, R> {
+	const pubsub = yield* PubSub.unbounded<A>({ replay: 1 });
+	const initial = yield* read();
+	PubSub.publishUnsafe(pubsub, initial);
+	const self = Object.create(Proto);
+	self.read = read;
+	self.write = write;
+	self.pubsub = pubsub;
+	self.semaphore = Semaphore.makeUnsafe(1);
+	return self;
+});
+
+/**
+ * Reads the current value.
+ */
+export const get = <A, E, R>(self: State<A, E, R>): Effect.Effect<A, E, R> =>
+	self.read();
+
+/**
+ * Replaces the value. Serialized with `update` / `modify` so writes
+ * happen in invocation order.
+ */
+export const set: {
+	<A>(value: A): <E, R>(self: State<A, E, R>) => Effect.Effect<void, E, R>;
+	<A, E, R>(self: State<A, E, R>, value: A): Effect.Effect<void, E, R>;
+} = dual(
+	2,
+	<A, E, R>(self: State<A, E, R>, value: A): Effect.Effect<void, E, R> =>
+		Semaphore.withPermit(self.semaphore, self.write(value)),
+);
+
+/**
+ * Updates the value by applying `f` to the current value. The
+ * read/apply/write triple is atomic across fibers.
+ */
+export const update: {
+	<A>(
+		f: (a: A) => A,
+	): <E, R>(self: State<A, E, R>) => Effect.Effect<void, E, R>;
+	<A, E, R>(self: State<A, E, R>, f: (a: A) => A): Effect.Effect<void, E, R>;
+} = dual(
+	2,
+	<A, E, R>(
+		self: State<A, E, R>,
+		f: (a: A) => A,
+	): Effect.Effect<void, E, R> =>
+		Semaphore.withPermit(
+			self.semaphore,
+			Effect.flatMap(self.read(), (a) => self.write(f(a))),
+		),
+);
+
+/**
+ * Updates the value by applying `f` and returns the new value. The
+ * read/apply/write triple is atomic across fibers.
+ */
+export const updateAndGet: {
+	<A>(f: (a: A) => A): <E, R>(self: State<A, E, R>) => Effect.Effect<A, E, R>;
+	<A, E, R>(self: State<A, E, R>, f: (a: A) => A): Effect.Effect<A, E, R>;
+} = dual(
+	2,
+	<A, E, R>(self: State<A, E, R>, f: (a: A) => A): Effect.Effect<A, E, R> =>
+		Semaphore.withPermit(
+			self.semaphore,
+			Effect.flatMap(self.read(), (a) => {
+				const next = f(a);
+				return Effect.as(self.write(next), next);
+			}),
+		),
+);
+
+/**
+ * Atomically replaces the value with the second element of `f(prev)`
+ * and returns the first. The read/apply/write triple is atomic across
+ * fibers.
+ */
+export const modify: {
+	<A, B>(
+		f: (a: A) => readonly [B, A],
+	): <E, R>(self: State<A, E, R>) => Effect.Effect<B, E, R>;
+	<A, E, R, B>(
+		self: State<A, E, R>,
+		f: (a: A) => readonly [B, A],
+	): Effect.Effect<B, E, R>;
+} = dual(
+	2,
+	<A, E, R, B>(
+		self: State<A, E, R>,
+		f: (a: A) => readonly [B, A],
+	): Effect.Effect<B, E, R> =>
+		Semaphore.withPermit(
+			self.semaphore,
+			Effect.flatMap(self.read(), (a) => {
+				const [b, next] = f(a);
+				return Effect.as(self.write(next), b);
+			}),
+		),
+);
+
+/**
+ * Stream of every value published to this `State`. New subscribers
+ * immediately see the most recent value (replay = 1), then every
+ * subsequent publish.
+ */
+export const changes = <A, E, R>(self: State<A, E, R>): Stream.Stream<A> =>
+	Stream.fromPubSub(self.pubsub);
+
+/**
+ * Publish a value to the change stream as an `Effect`. Does not
+ * modify the underlying store.
+ */
+export const publish: {
+	<A>(value: A): <E, R>(self: State<A, E, R>) => Effect.Effect<boolean>;
+	<A, E, R>(self: State<A, E, R>, value: A): Effect.Effect<boolean>;
+} = dual(
+	2,
+	<A, E, R>(self: State<A, E, R>, value: A): Effect.Effect<boolean> =>
+		PubSub.publish(self.pubsub, value),
+);
+
+/**
+ * Synchronous variant of {@link publish}. Returns `true` when the
+ * publish succeeded, `false` if the pubsub is shut down. The runtime
+ * uses this from rivetkit's `onStateChange` callback to feed the
+ * change stream.
+ */
+export const publishUnsafe = <A, E, R>(
+	self: State<A, E, R>,
+	value: A,
+): boolean => PubSub.publishUnsafe(self.pubsub, value);

package/src/internal/ActionDispatcher.ts

@@ -0,0 +1,192 @@
+import {
+	Cause,
+	Effect,
+	Exit,
+	type Fiber,
+	Option,
+	Record,
+	Schema,
+	Tracer,
+} from "effect";
+import * as Rivetkit from "rivetkit";
+import type * as Action from "../Action.ts";
+import type {
+	ActionHandlersFrom,
+	ActionRequest,
+	Actor,
+} from "../Actor.ts";
+import type * as Client from "../Client.ts";
+import * as ActionErrorEnvelope from "./ActionErrorEnvelope.ts";
+import { makeActorLogAnnotations } from "./logging.ts";
+import { readTraceMeta, rpcSystem } from "./tracing.ts";
+import { hasStringProperty } from "./utils.ts";
+
+export type Instance<ActionHandlers> = {
+	readonly actionHandlers: ActionHandlers;
+	readonly runFork: <A, E>(
+		effect: Effect.Effect<A, E, any>,
+		options?: Effect.RunOptions,
+	) => Fiber.Fiber<A, E>;
+};
+
+export const make = <
+	Name extends string,
+	Actions extends Action.AnyWithProps,
+	ActionHandlers extends ActionHandlersFrom<Actions>,
+	ActorDefinition extends Rivetkit.AnyActorDefinition,
+>({
+	actor,
+	getInstance,
+}: {
+	readonly actor: Actor<Name, Actions>;
+	readonly getInstance: (
+		actorId: string,
+	) => Instance<ActionHandlers> | undefined;
+}) =>
+	Record.fromIterableWith(actor.actions, (action) => {
+		const decodePayload = Schema.decodeUnknownEffect(
+			Schema.toCodecJson(action.payloadSchema),
+		);
+		const encodeSuccess = Schema.encodeEffect(
+			Schema.toCodecJson(action.successSchema),
+		);
+		const encodeError = Schema.encodeEffect(
+			Schema.toCodecJson(action.errorSchema),
+		);
+
+		return [
+			action._tag,
+			async (
+				c: Rivetkit.ActionContextOf<ActorDefinition>,
+				payload: Action.Payload<typeof action>,
+				meta?: Client.ActionMeta, // TODO: Find better type
+			) => {
+				// Always wrap in a server-side span so the handler has a
+				// live `currentSpan` even when the caller didn't ship trace
+				// context (e.g., a non-Effect-SDK client). When trace context
+				// is present, reattach it as the parent so the server span
+				// joins the caller's trace.
+				const rpcMethod = `${actor.name}/${action._tag}`;
+				const traceMeta = readTraceMeta(meta);
+
+				const instance = getInstance(c.actorId);
+				if (!instance) {
+					if (c.abortSignal.aborted) throw makeActorAbortedError();
+					throw new Error("actor instance missing");
+				}
+
+				const actionEffect = Effect.gen(function* () {
+					// The handler map is keyed by the same action
+					// definitions being registered here, but
+					// TypeScript loses that relationship once the
+					// actions are widened into the RivetKit actions
+					// record.
+					const actionHandler = instance.actionHandlers[
+						action._tag as keyof ActionHandlers
+					] as (
+						envelope: ActionRequest<typeof action>,
+					) => Action.ResultFrom<typeof action, any>;
+					// Raw RivetKit clients call no-argument actions with an
+					// absent first argument. The Effect JSON Void codec expects
+					// null, so adapt only actions that declared no payload.
+					const payloadForDecode =
+						!action.hasPayload && payload === undefined
+							? null
+							: payload;
+					const decodedPayload = yield* decodePayload(
+						payloadForDecode,
+					).pipe(
+						Effect.mapError(() =>
+							new Rivetkit.RivetError(
+								"request",
+								"invalid",
+								`Invalid payload for action ${actor.name}/${action._tag}`,
+							),
+						),
+					);
+					// The payload was decoded with this action's schema,
+					// so this is the runtime boundary that restores the
+					// typed envelope expected by the user handler.
+					const actionRequest = {
+						_tag: action._tag,
+						action,
+						payload: decodedPayload,
+					} as ActionRequest<typeof action>;
+
+					const resultExit = yield* Effect.exit(
+						actionHandler(actionRequest),
+					);
+
+					if (Exit.isSuccess(resultExit)) {
+						return yield* encodeSuccess(resultExit.value).pipe(
+							Effect.orDie,
+						);
+					}
+
+					const expectedError = Exit.findErrorOption(resultExit);
+
+					if (Option.isSome(expectedError)) {
+						const encodedError = yield* encodeError(
+							expectedError.value,
+						).pipe(Effect.orDie);
+
+						return yield* Effect.fail(
+							new Rivetkit.UserError(
+								hasStringProperty("message")(encodedError)
+									? encodedError.message
+									: `${action._tag} failed`,
+								{
+									code: hasStringProperty("_tag")(
+										encodedError,
+									)
+										? encodedError._tag
+										: undefined,
+									metadata:
+										ActionErrorEnvelope.make(encodedError),
+								},
+							),
+						);
+					}
+
+					return yield* Effect.failCause(resultExit.cause);
+				}).pipe(
+					Effect.withSpan(rpcMethod, {
+						parent: traceMeta
+							? Tracer.externalSpan(traceMeta)
+							: undefined,
+						kind: "server",
+						attributes: {
+							"rpc.system.name": rpcSystem,
+							"rpc.method": rpcMethod,
+						},
+					}),
+					Effect.annotateLogs(makeActorLogAnnotations(c)),
+				);
+				const fiber = instance.runFork(actionEffect, {
+					signal: c.abortSignal,
+				});
+				const exit = await new Promise<Exit.Exit<unknown, unknown>>(
+					(resolve) => fiber.addObserver(resolve),
+				);
+
+				if (Exit.isSuccess(exit)) return exit.value;
+				// Action fibers can be interrupted by a caller abort signal
+				// or by the actor instance scope closing during sleep, destroy,
+				// or shutdown. Surface those lifecycle exits as RivetKit's
+				// structured action-aborted error instead of an internal error.
+				if (Cause.hasInterruptsOnly(exit.cause)) {
+					throw makeActorAbortedError();
+				}
+				const expectedError = Exit.findErrorOption(exit);
+				if (Option.isSome(expectedError)) {
+					throw expectedError.value;
+				}
+				throw Cause.squash(exit.cause);
+			},
+		];
+	});
+
+const makeActorAbortedError = () =>
+	new Rivetkit.RivetError("actor", "aborted", "Actor aborted", {
+		public: true,
+	});

package/src/internal/ActionErrorEnvelope.ts

@@ -0,0 +1,19 @@
+import { Schema } from "effect";
+
+export const tag = "EffectActionError" as const;
+
+export const schemaVersion = 1 as const;
+
+export const ActionErrorEnvelope = Schema.Struct({
+	_tag: Schema.tag(tag),
+	version: Schema.Literal(schemaVersion),
+	error: Schema.Unknown,
+});
+
+export type ActionErrorEnvelope = typeof ActionErrorEnvelope.Type;
+
+export const make = (error: unknown): ActionErrorEnvelope => ({
+	_tag: tag,
+	version: schemaVersion,
+	error,
+});