fix: preserve task and entity type safety

.changeset/honest-task-entity-types.md

@@ -0,0 +1,5 @@
+---
+"@humanlayer/effect-machine": patch
+---
+
+Preserve exact machine schema types on `MachineEntity` and require `.task()` shorthand to return a valid machine event without using `any` or diagnostic suppressions.

src/cluster/entity-actor-ref.ts

@@ -15,6 +15,7 @@ import { Effect, Option, Schema, Stream } from "effect";
 
 import type { ExtractReply, ReplyTypeBrand } from "../internal/brands.js";
 import { ActorStoppedError, NoReplyError } from "../errors.js";
+import type { MachineSchemaDefinition } from "../schema.js";
 import type { EntityRpcs, MachineEntity } from "./to-entity.js";
 
 /**
@@ -75,8 +76,10 @@ export const makeEntityActorRef = <
   R,
   EntityType extends string,
   ClientError,
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
 >(
-  entity: MachineEntity<State, Event, R, EntityType>,
+  entity: MachineEntity<State, Event, R, EntityType, StateDefinition, EventDefinition>,
   client: RpcClient.RpcClient<
     EntityRpcs<Schema.Codec<State>, Schema.Codec<Event>>[number],
     ClientError

src/cluster/entity-machine.ts

@@ -29,6 +29,7 @@ import {
 } from "effect";
 
 import { type Machine, replay } from "../machine.js";
+import type { MachineSchemaDefinition } from "../schema.js";
 import type { ActorSystemService } from "../actor.js";
 import { ActorSystem as ActorSystemTag, makeSystem } from "../actor.js";
 import type { ProcessEventHooks } from "../internal/transition.js";
@@ -139,34 +140,45 @@ function layer<
   E extends { readonly _tag: string },
   R,
   EntityType extends string,
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
 >(
-  entity: MachineEntity<S, E, R, EntityType>,
+  entity: MachineEntity<S, E, R, EntityType, StateDefinition, EventDefinition>,
   options?: EntityOptionsWithoutPersistence<S, E>,
 ): EntityLayer<R, never>;
 function layer<
   S extends { readonly _tag: string },
   E extends { readonly _tag: string },
   R,
   EntityType extends string,
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
 >(
-  entity: MachineEntity<S, E, R, EntityType>,
+  entity: MachineEntity<S, E, R, EntityType, StateDefinition, EventDefinition>,
   options: EntityOptionsWithPersistence<S, E>,
 ): EntityLayer<R, PersistenceAdapter>;
 function layer<
   S extends { readonly _tag: string },
   E extends { readonly _tag: string },
   R,
   EntityType extends string,
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
 >(
-  entity: MachineEntity<S, E, R, EntityType>,
+  entity: MachineEntity<S, E, R, EntityType, StateDefinition, EventDefinition>,
   options: EntityMachineOptions<S, E>,
 ): EntityLayer<R, PersistenceAdapter>;
 function layer<
   S extends { readonly _tag: string },
   E extends { readonly _tag: string },
   R,
   EntityType extends string,
->(entity: MachineEntity<S, E, R, EntityType>, options?: EntityMachineOptions<S, E>) {
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
+>(
+  entity: MachineEntity<S, E, R, EntityType, StateDefinition, EventDefinition>,
+  options?: EntityMachineOptions<S, E>,
+) {
   type Rpcs = EntityRpcs<Schema.Codec<S>, Schema.Codec<E>>[number];
   const machine = entity.machine;
   const persistence = options?.persistence;

src/cluster/to-entity.ts

@@ -8,6 +8,7 @@ import { Rpc } from "effect/unstable/rpc";
 import { Schema } from "effect";
 
 import type { Machine } from "../machine.js";
+import type { MachineSchemaDefinition } from "../schema.js";
 import { MissingSchemaError } from "../errors.js";
 
 /**
@@ -60,9 +61,13 @@ export interface MachineEntity<
   Event extends { readonly _tag: string },
   R,
   EntityType extends string,
+  StateDefinition extends Record<string, Schema.Struct.Fields> = Record<
+    string,
+    Schema.Struct.Fields
+  >,
+  EventDefinition extends MachineSchemaDefinition = MachineSchemaDefinition,
 > extends Entity.Entity<EntityType, EntityRpcs<Schema.Codec<State>, Schema.Codec<Event>>[number]> {
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any -- schema-definition parameters are carried opaquely by MachineEntity
-  readonly machine: Machine<State, Event, R, any, any>;
+  readonly machine: Machine<State, Event, R, StateDefinition, EventDefinition>;
   readonly stateSchema: Schema.Codec<State>;
   readonly eventSchema: Schema.Codec<Event>;
   readonly rpcs: EntityRpcs<Schema.Codec<State>, Schema.Codec<Event>>;
@@ -104,11 +109,12 @@ export const toEntity = <
   E extends { readonly _tag: string },
   R,
   const EntityType extends string,
+  StateDefinition extends Record<string, Schema.Struct.Fields>,
+  EventDefinition extends MachineSchemaDefinition,
 >(
-  // eslint-disable-next-line @typescript-eslint/no-explicit-any -- Schema fields need wide acceptance
-  machine: Machine<S, E, R, any, any>,
+  machine: Machine<S, E, R, StateDefinition, EventDefinition>,
   options: ToEntityOptions<EntityType>,
-): MachineEntity<S, E, R, EntityType> => {
+): MachineEntity<S, E, R, EntityType, StateDefinition, EventDefinition> => {
   const stateSchema = machine.stateSchema;
   const eventSchema = machine.eventSchema;
 

src/machine.ts

@@ -34,8 +34,8 @@
  *
  * @module
  */
-import type { Duration, Schema } from "effect";
-import { Cause, Effect, Exit, Option, Random, Scope } from "effect";
+import type { Duration } from "effect";
+import { Cause, Effect, Exit, Option, Random, Schema, Scope } from "effect";
 
 import type { DeferReplyResult, ReplyResult, TransitionResult } from "./internal/utils.js";
 import { getTag, makeReply, makeDeferReply } from "./internal/utils.js";
@@ -178,12 +178,18 @@ export interface BackgroundEffect<State, Event, R> {
 // Options types
 // ============================================================================
 
-export interface TaskOptions<State, Event, A, E1, ES, EF> {
-  readonly onSuccess?: (value: A, ctx: StateHandlerContext<State, Event>) => ES;
+interface TaskCommonOptions<State, Event, E1, EF> {
   readonly onFailure?: (cause: Cause.Cause<E1>, ctx: StateHandlerContext<State, Event>) => EF;
   readonly name?: string;
 }
 
+/** Options for a mapped task, or for shorthand when the task already returns an event. */
+export type TaskOptions<State, Event, A, E1, ES, EF> = TaskCommonOptions<State, Event, E1, EF> &
+  (
+    | { readonly onSuccess: (value: A, ctx: StateHandlerContext<State, Event>) => ES }
+    | ([A] extends [Event] ? { readonly onSuccess?: undefined } : never)
+  );
+
 // ============================================================================
 // Recovery / Durability
 // ============================================================================
@@ -633,58 +639,58 @@ export class Machine<
     A,
     E1,
     R1,
-    ES extends VariantsUnion<_ED> & BrandedEvent,
-    EF extends VariantsUnion<_ED> & BrandedEvent,
+    ES extends Event & VariantsUnion<_ED> & BrandedEvent,
+    EF extends Event & VariantsUnion<_ED> & BrandedEvent,
   >(
     state: TaggedOrConstructor<NS>,
-    run: (
-      ctx: StateHandlerContext<NS, VariantsUnion<_ED> & BrandedEvent>,
-    ) => Effect.Effect<A, E1, Scope.Scope | R1>,
-    options: TaskOptions<NS, VariantsUnion<_ED> & BrandedEvent, A, E1, ES, EF>,
+    run: (ctx: StateHandlerContext<NS, Event>) => Effect.Effect<A, E1, Scope.Scope | R1>,
+    options: TaskOptions<NS, Event, A, E1, ES, EF>,
   ): Machine<State, Event, R | R1, _SD, _ED>;
   /** Multiple states — handler receives the selected state union. */
   task<
     NS extends ReadonlyArray<TaggedOrConstructor<VariantsUnion<_SD> & BrandedState>>,
     A,
     E1,
     R1,
-    ES extends VariantsUnion<_ED> & BrandedEvent,
-    EF extends VariantsUnion<_ED> & BrandedEvent,
+    ES extends Event & VariantsUnion<_ED> & BrandedEvent,
+    EF extends Event & VariantsUnion<_ED> & BrandedEvent,
   >(
     states: NS,
     run: (
-      ctx: StateHandlerContext<
-        NS[number] extends TaggedOrConstructor<infer S> ? S : never,
-        VariantsUnion<_ED> & BrandedEvent
-      >,
+      ctx: StateHandlerContext<NS[number] extends TaggedOrConstructor<infer S> ? S : never, Event>,
     ) => Effect.Effect<A, E1, Scope.Scope | R1>,
     options: TaskOptions<
       NS[number] extends TaggedOrConstructor<infer S> ? S : never,
-      VariantsUnion<_ED> & BrandedEvent,
+      Event,
       A,
       E1,
       ES,
       EF
     >,
   ): Machine<State, Event, R | R1, _SD, _ED>;
-  /* eslint-disable @typescript-eslint/no-explicit-any -- public overloads preserve selection/task correlation at this implementation boundary */
-  task(
-    stateOrStates: any,
-    run: (ctx: StateHandlerContext<any, any>) => Effect.Effect<any, any, any>,
-    options: any,
-  ): Machine<State, Event, R, _SD, _ED> {
-    const handler: StateEffectHandler<any, any, any> = Effect.fn("effect-machine.task")(function* (
-      // eslint-disable-next-line @typescript-eslint/no-explicit-any -- implementation is checked by public overloads
-      ctx: StateHandlerContext<any, any>,
-    ) {
+  task<
+    NS extends VariantsUnion<_SD> & BrandedState,
+    A,
+    E1,
+    R1,
+    ES extends Event & VariantsUnion<_ED> & BrandedEvent,
+    EF extends Event & VariantsUnion<_ED> & BrandedEvent,
+  >(
+    stateOrStates: TaggedOrConstructor<NS> | ReadonlyArray<TaggedOrConstructor<NS>>,
+    run: (ctx: StateHandlerContext<NS, Event>) => Effect.Effect<A, E1, Scope.Scope | R1>,
+    options: TaskOptions<NS, Event, A, E1, ES, EF>,
+  ): Machine<State, Event, R | R1, _SD, _ED> {
+    const isEvent = this.eventSchema !== undefined ? Schema.is(this.eventSchema) : undefined;
+    const handler: StateEffectHandler<NS, Event, Scope.Scope | R1> = Effect.fn(
+      "effect-machine.task",
+    )(function* (ctx: StateHandlerContext<NS, Event>) {
       yield* emitTaskInspection({
         actorId: ctx.actorId,
         state: ctx.state,
         taskName: options.name,
         phase: "start",
       });
 
-      // @effect-diagnostics anyUnknownInErrorContext:off -- the public task overloads preserve concrete error and requirement channels at this implementation boundary
       const exit = yield* Effect.exit(run(ctx));
 
       if (Exit.isSuccess(exit)) {
@@ -695,7 +701,11 @@ export class Machine<
           phase: "success",
         });
         const successEvent =
-          options.onSuccess !== undefined ? options.onSuccess(exit.value, ctx) : exit.value;
+          options.onSuccess !== undefined
+            ? options.onSuccess(exit.value, ctx)
+            : isEvent !== undefined && isEvent(exit.value)
+              ? exit.value
+              : yield* Effect.die("Task shorthand produced an invalid machine event");
         yield* ctx.self.send(successEvent);
         yield* Effect.yieldNow;
         return;
@@ -723,13 +733,11 @@ export class Machine<
         yield* Effect.yieldNow;
         return;
       }
-      // @effect-diagnostics anyUnknownInErrorContext:off
       return yield* Effect.failCause(cause).pipe(Effect.orDie);
     });
 
     return this.registerStateEffect(stateOrStates, handler);
   }
-  /* eslint-enable @typescript-eslint/no-explicit-any */
 
   // ---- timeout ----
 
@@ -756,7 +764,7 @@ export class Machine<
    */
   timeout<NS extends VariantsUnion<_SD> & BrandedState>(
     state: TaggedOrConstructor<NS>,
-    config: TimeoutConfig<NS, VariantsUnion<_ED> & BrandedEvent>,
+    config: TimeoutConfig<NS, Event & VariantsUnion<_ED> & BrandedEvent>,
   ): Machine<State, Event, R, _SD, _ED> {
     const stateTag = getTag(state);
     const duration = config.duration;

test/cluster-type-constraints.test.ts

@@ -32,6 +32,14 @@ const clusterMachine = Machine.make({
   .background(() => ClusterService.pipe(Effect.andThen((service) => service.run)));
 
 const ClusterEntity = toEntity(clusterMachine, { type: "TypeConstraints" });
+
+const ForeignState = State({ Foreign: {} });
+const ForeignEvent = Event({ Foreign: {} });
+// @ts-expect-error - a machine-owned entity retains its machine's exact state/event definitions
+ClusterEntity.machine.on(ForeignState.Foreign, ForeignEvent.Foreign, () =>
+  ClusterState.Active({ count: 0 }),
+);
+
 const withoutPersistence = EntityMachine.layer(ClusterEntity);
 const withPersistence = EntityMachine.layer(ClusterEntity, {
   persistence: { strategy: "journal" },

test/type-constraints.test.ts

@@ -104,6 +104,14 @@ const _test6 = Machine.make({
   onSuccess: () => MyEvent.Complete,
 });
 
+// Task shorthand is only valid when the task itself returns a machine event.
+const _invalidTaskShorthand = Machine.make({
+  state: MyState,
+  event: MyEvent,
+  initial: MyState.Loading({ url: "/" }),
+  // @ts-expect-error - non-event task results require an onSuccess mapper
+}).task(MyState.Loading, () => Effect.succeed(123), {});
+
 const _test6Spawn = Machine.spawn(_test6);
 type _Test6RequiresService = Assert<
   MyService extends EffectRequirements<typeof _test6Spawn> ? true : false