@shivaedev/effect-contract 0.11.0 → 0.13.0

package/src/bind.ts

@@ -1,17 +1,17 @@
-import type { Duration, Schema } from "effect";
-import { Effect, Stream } from "effect";
+import type { Duration, Effect, Stream } from "effect";
 import type { Headers } from "effect/http";
 import type * as AsyncResult from "effect/reactivity/AsyncResult";
 import type * as Atom from "effect/reactivity/Atom";
+import type * as AtomRegistry from "effect/reactivity/AtomRegistry";
 import type * as AtomRpc from "effect/reactivity/AtomRpc";
-import * as Reactivity from "effect/reactivity/Reactivity";
-import type { Rpc, RpcSchema } from "effect/rpc";
+import type * as Reactivity from "effect/reactivity/Reactivity";
+import type { Rpc } from "effect/rpc";
 import type { RpcClientError } from "effect/rpc/RpcClientError";
-import { queryReactivity } from "#internal/queryReactivity.ts";
-import { automaticLive, type LiveBound, LiveConnection } from "./automaticLive.ts";
-import type { Contract, Declared, Tag } from "./contract.ts";
-import { invalidationKeys, type Key, readKeys } from "./keys.ts";
-import type { CommandShape, QueryShape, StreamShape } from "./operation.ts";
+import { automaticLive, type LiveBound, LiveConnection } from "#automaticLive.ts";
+import type { Contract, Declared, Tag } from "#contract.ts";
+import { bindingOwner } from "#internal/bindingOwner.ts";
+import { boundCommand, boundQuery, boundStream } from "#internal/clientOperations.ts";
+import type { CommandShape, QueryShape, StreamShape } from "#operation.ts";
 
 export type Failure<R extends Rpc.Any> =
 	R extends Rpc.Rpc<infer _Tag, infer _Payload, infer _Success, infer _Error, infer TMiddleware, infer _Requires>
@@ -57,51 +57,12 @@ export type Bound<
 	readonly [TStream in TStreams[number] as TStream["name"]]: BoundStream<Rpc.ExtractTag<TRpcs, Tag<TName, TStream["name"]>>, TSelf>;
 };
 
-type ErasedClient = AtomRpc.AtomRpcClient<unknown, string, Rpc.Rpc<string, Schema.Top, Schema.Top, Schema.Top>>;
-
-type ErasedStreamingClient = AtomRpc.AtomRpcClient<
-	unknown,
-	string,
-	Rpc.Rpc<string, Schema.Top, RpcSchema.Stream<Schema.Top, Schema.Top>, Schema.Never>
->;
-
 type MatchingClient<TRpcs extends Rpc.Any, TClientRpcs extends Rpc.Any> = [TRpcs] extends [Extract<TClientRpcs, { readonly _tag: TRpcs["_tag"] }>]
 	? [Extract<TClientRpcs, { readonly _tag: TRpcs["_tag"] }>] extends [TRpcs]
 		? unknown
 		: "Client must preserve the contract RPC types"
 	: "Client must implement the contract RPCs";
 
-const invalidateAfter = Effect.fn("EffectContract.invalidateAfter")(function* (keys: readonly Key[]) {
-	yield* Reactivity.invalidate(invalidationKeys(keys));
-});
-
-function boundQuery(service: ErasedClient, tag: string, query: QueryShape, track: (keys: readonly Key[]) => void) {
-	const reactive = queryReactivity(service.runtime);
-	function reads(payload: unknown) {
-		return readKeys(query.reads(payload));
-	}
-	return {
-		query: (payload: unknown, options: QueryOptions = {}) => {
-			track(query.reads(payload));
-			return reactive(service.query(tag, payload, options), reads(payload));
-		},
-		run: (payload: unknown) => service.use((client) => client(tag, payload)),
-	};
-}
-
-function boundCommand(service: ErasedClient, tag: string, command: CommandShape) {
-	return {
-		run: (payload: unknown) =>
-			service.use((client) => client(tag, payload)).pipe(Effect.tap((result) => invalidateAfter(command.invalidates(payload, result)))),
-	};
-}
-
-function boundStream(service: ErasedStreamingClient, tag: string) {
-	return {
-		run: (payload: unknown) => Stream.unwrap(service.use((client) => Effect.succeed(client(tag, payload)))),
-	};
-}
-
 export function bind<
 	TName extends string,
 	TQueries extends readonly QueryShape[],
@@ -111,20 +72,26 @@ export function bind<
 	TId extends string,
 	TClientRpcs extends Rpc.Any,
 	TStreams extends readonly StreamShape[],
+	TScope extends "session" | undefined,
 >(
-	contract: Contract<TName, TQueries, TCommands, TRpcs, TStreams>,
+	contract: Contract<TName, TQueries, TCommands, TRpcs, TStreams, TScope>,
 	service: AtomRpc.AtomRpcClient<TSelf, TId, TClientRpcs> & MatchingClient<NoInfer<TRpcs>, TClientRpcs>,
+	...options: "session" extends TScope
+		? [options: { readonly registry: AtomRegistry.AtomRegistry }]
+		: [options?: { readonly registry: AtomRegistry.AtomRegistry }]
 ): Bound<TName, TQueries, TCommands, TRpcs, TSelf, TStreams> & LiveBound;
 export function bind(
-	contract: { readonly declaration: Declared<string, readonly QueryShape[], readonly CommandShape[], readonly StreamShape[]> },
+	contract: { readonly declaration: Declared<string, readonly QueryShape[], readonly CommandShape[], readonly StreamShape[], "session" | undefined> },
 	service: AtomRpc.AtomRpcClient<unknown, string, Rpc.Any>,
+	options?: { readonly registry: AtomRegistry.AtomRegistry },
 ): unknown {
-	const { name, queries, commands, streams } = contract.declaration;
-	const automatic = automaticLive(service, name);
+	const { name, queries, commands, streams, keyScope } = contract.declaration;
+	const owner = bindingOwner(keyScope, options?.registry);
+	const automatic = automaticLive(service, name, owner);
 	return Object.assign(
 		Object.fromEntries([
-			...queries.map((query) => [query.name, boundQuery(service, `${name}.${query.name}`, query, automatic.track)] as const),
-			...commands.map((command) => [command.name, boundCommand(service, `${name}.${command.name}`, command)] as const),
+			...queries.map((query) => [query.name, boundQuery(service, `${name}.${query.name}`, query, automatic.track, owner)] as const),
+			...commands.map((command) => [command.name, boundCommand(service, `${name}.${command.name}`, command, owner)] as const),
 			...streams.map((stream) => [stream.name, boundStream(service, `${name}.${stream.name}`)] as const),
 		]),
 		{ [LiveConnection]: automatic.subscription },

package/src/contract.ts

@@ -1,7 +1,7 @@
 import { Context, type Schema } from "effect";
 import { Rpc, RpcGroup, type RpcMiddleware, type RpcSchema } from "effect/rpc";
-import type { CommandShape, OperationShape, QueryShape, StreamShape } from "./operation.ts";
-import { CommandInvalidation, changesRpc, QueryReads } from "./serverMetadata.ts";
+import type { CommandShape, OperationShape, QueryShape, StreamShape } from "#operation.ts";
+import { CommandInvalidation, changesRpc, KeyScope, QueryReads } from "#serverMetadata.ts";
 
 export type Tag<TContract extends string, TName extends string> = `${TContract}.${TName}`;
 
@@ -16,8 +16,10 @@ export interface Declared<
 	TQueries extends readonly QueryShape[],
 	TCommands extends readonly CommandShape[],
 	TStreams extends readonly StreamShape[] = readonly [],
+	TScope extends "session" | undefined = undefined,
 > {
 	readonly commands: TCommands;
+	readonly keyScope: TScope;
 	readonly name: TName;
 	readonly queries: TQueries;
 	readonly streams: TStreams;
@@ -27,12 +29,15 @@ export interface Contract<
 	TName extends string,
 	TQueries extends readonly QueryShape[],
 	TCommands extends readonly CommandShape[],
-	TRpcs extends Rpc.Any = OperationRpc<TName, TQueries[number] | TCommands[number]> | ReturnType<typeof changesRpc<TName>>,
+	TRpcs extends Rpc.Any = OperationRpc<TName, TQueries[number] | TCommands[number]> | ReturnType<typeof changesRpc<TName, TQueries>>,
 	TStreams extends readonly StreamShape[] = readonly [],
+	TScope extends "session" | undefined = undefined,
 > extends RpcGroup.RpcGroup<TRpcs> {
 	readonly application: RpcGroup.RpcGroup<Exclude<TRpcs, { readonly _tag: `${TName}.$changes` }>>;
-	readonly declaration: Declared<TName, TQueries, TCommands, TStreams>;
-	middleware: <TM extends RpcMiddleware.AnyService>(middleware: TM) => Contract<TName, TQueries, TCommands, Rpc.AddMiddleware<TRpcs, TM>, TStreams>;
+	readonly declaration: Declared<TName, TQueries, TCommands, TStreams, TScope>;
+	middleware: <TM extends RpcMiddleware.AnyService>(
+		middleware: TM,
+	) => Contract<TName, TQueries, TCommands, Rpc.AddMiddleware<TRpcs, TM>, TStreams, TScope>;
 }
 
 type Duplicated<TOperations extends readonly OperationShape[], TSeen extends string = never> = TOperations extends readonly [
@@ -53,14 +58,15 @@ function assertUnique(operations: readonly OperationShape[]) {
 	}
 }
 
-function annotations(operation: QueryShape | CommandShape | StreamShape) {
+function annotations(operation: QueryShape | CommandShape | StreamShape, keyScope: "session" | undefined) {
+	const scope = keyScope === undefined ? Context.empty() : Context.make(KeyScope, keyScope);
 	if (operation.kind === "command") {
-		return Context.make(CommandInvalidation, operation.invalidates);
+		return Context.add(scope, CommandInvalidation, operation.invalidates);
 	}
-	return operation.kind === "query" ? Context.make(QueryReads, operation.reads) : Context.empty();
+	return operation.kind === "query" ? Context.add(scope, QueryReads, operation.reads) : scope;
 }
 
-type AnyDeclared = Declared<string, readonly QueryShape[], readonly CommandShape[], readonly StreamShape[]>;
+type AnyDeclared = Declared<string, readonly QueryShape[], readonly CommandShape[], readonly StreamShape[], "session" | undefined>;
 
 interface Group {
 	middleware: (middleware: RpcMiddleware.AnyService) => Group;
@@ -86,9 +92,11 @@ export function contract<
 	const TQueries extends readonly QueryShape[] = readonly [],
 	const TCommands extends readonly CommandShape[] = readonly [],
 	const TStreams extends readonly StreamShape[] = readonly [],
+	const TScope extends "session" | undefined = undefined,
 >(
 	name: TName,
 	operations: {
+		readonly keyScope?: TScope;
 		readonly queries?: TQueries & Unique<TQueries>;
 		readonly commands?: TCommands & Unique<[...TQueries, ...TCommands]>;
 		readonly streams?: TStreams & Unique<[...TQueries, ...TCommands, ...TStreams]>;
@@ -97,14 +105,20 @@ export function contract<
 	TName,
 	TQueries,
 	TCommands,
-	OperationRpc<TName, TQueries[number] | TCommands[number] | TStreams[number]> | ReturnType<typeof changesRpc<TName>>,
-	TStreams
+	OperationRpc<TName, TQueries[number] | TCommands[number] | TStreams[number]> | ReturnType<typeof changesRpc<TName, TQueries>>,
+	TStreams,
+	TScope
 >;
 export function contract(
 	name: string,
-	operations: { readonly queries?: readonly QueryShape[]; readonly commands?: readonly CommandShape[]; readonly streams?: readonly StreamShape[] },
+	operations: {
+		readonly keyScope?: "session";
+		readonly queries?: readonly QueryShape[];
+		readonly commands?: readonly CommandShape[];
+		readonly streams?: readonly StreamShape[];
+	},
 ): DeclaredGroup {
-	const { queries = [], commands = [], streams = [] } = operations;
+	const { queries = [], commands = [], streams = [], keyScope } = operations;
 	assertUnique([...queries, ...commands, ...streams]);
 	const rpcs = [...queries, ...commands, ...streams].map((operation) =>
 		Rpc.make(`${name}.${operation.name}`, {
@@ -112,8 +126,15 @@ export function contract(
 			payload: operation.payload,
 			stream: operation.kind === "stream",
 			success: operation.success,
-		}).annotateMerge(annotations(operation)),
+		}).annotateMerge(annotations(operation, keyScope)),
 	);
 	const application = RpcGroup.make(...rpcs);
-	return withDeclaration(application.add(changesRpc(name)), application, { commands, name, queries, streams });
+	const changes = changesRpc(name, queries);
+	return withDeclaration(application.add(keyScope === undefined ? changes : changes.annotate(KeyScope, keyScope)), application, {
+		commands,
+		keyScope,
+		name,
+		queries,
+		streams,
+	});
 }

package/src/internal/bindingOwner.ts

@@ -0,0 +1,24 @@
+import type * as AtomRegistry from "effect/reactivity/AtomRegistry";
+import type { KeyContext } from "#keys.ts";
+
+export interface BindingOwner {
+	readonly context: KeyContext;
+	readonly registry: AtomRegistry.AtomRegistry;
+}
+
+const contexts = new WeakMap<AtomRegistry.AtomRegistry, KeyContext>();
+
+export function bindingOwner(scope: "session" | undefined, registry?: AtomRegistry.AtomRegistry): BindingOwner | undefined {
+	if (scope === undefined) {
+		return undefined;
+	}
+	if (registry === undefined) {
+		throw new Error("A session-scoped contract requires its owning registry");
+	}
+	let context = contexts.get(registry);
+	if (context === undefined) {
+		context = { scope: crypto.randomUUID() };
+		contexts.set(registry, context);
+	}
+	return { context, registry };
+}

package/src/internal/clientOperations.ts

@@ -0,0 +1,74 @@
+import { Effect, type Schema, Stream } from "effect";
+import * as AtomRegistry from "effect/reactivity/AtomRegistry";
+import type * as AtomRpc from "effect/reactivity/AtomRpc";
+import * as Reactivity from "effect/reactivity/Reactivity";
+import type { Rpc, RpcSchema } from "effect/rpc";
+import type { automaticLive } from "#automaticLive.ts";
+import type { QueryOptions } from "#bind.ts";
+import { invalidationKeys, readKeys } from "#keys.ts";
+import type { CommandShape, QueryShape } from "#operation.ts";
+import type { BindingOwner } from "./bindingOwner.ts";
+import { queryReactivity } from "./queryReactivity.ts";
+
+type ErasedClient = AtomRpc.AtomRpcClient<unknown, string, Rpc.Rpc<string, Schema.Top, Schema.Top, Schema.Top>>;
+
+type ErasedStreamingClient = AtomRpc.AtomRpcClient<
+	unknown,
+	string,
+	Rpc.Rpc<string, Schema.Top, RpcSchema.Stream<Schema.Top, Schema.Top>, Schema.Never>
+>;
+
+function run<A, E>(service: ErasedClient, effect: Effect.Effect<A, E, unknown>, owner?: BindingOwner) {
+	return owner === undefined ? effect : Effect.suspend(() => AtomRegistry.getResult(owner.registry, service.runtime.atom(effect)));
+}
+
+export function boundQuery(
+	service: ErasedClient,
+	tag: string,
+	query: QueryShape,
+	track: ReturnType<typeof automaticLive>["track"],
+	owner?: BindingOwner,
+) {
+	const reactive = queryReactivity(service.runtime);
+	return {
+		query: (payload: unknown, options: QueryOptions = {}) => {
+			const atom = service.query(tag, payload, options);
+			const keys = query.reads(payload);
+			const reference = { payload, query: query.name };
+			if (owner === undefined) {
+				track(undefined, reference, keys);
+			}
+			return reactive(
+				atom,
+				readKeys(keys, owner?.context),
+				(registry) => (owner === undefined ? () => undefined : track(registry, reference, keys)),
+				owner?.registry,
+			);
+		},
+		run: (payload: unknown) =>
+			run(
+				service,
+				service.use((client) => client(tag, payload)),
+				owner,
+			),
+	};
+}
+
+export function boundCommand(service: ErasedClient, tag: string, command: CommandShape, owner?: BindingOwner) {
+	return {
+		run: (payload: unknown) =>
+			run(
+				service,
+				service
+					.use((client) => client(tag, payload))
+					.pipe(Effect.tap((result) => Reactivity.invalidate(invalidationKeys(command.invalidates(payload, result), owner?.context)))),
+				owner,
+			),
+	};
+}
+
+export function boundStream(service: ErasedStreamingClient, tag: string) {
+	return {
+		run: (payload: unknown) => Stream.unwrap(service.use((client) => Effect.succeed(client(tag, payload)))),
+	};
+}

package/src/internal/queryReactivity.ts

@@ -1,6 +1,7 @@
 import { Duration, Effect } from "effect";
 import * as AsyncResult from "effect/reactivity/AsyncResult";
 import * as Atom from "effect/reactivity/Atom";
+import type * as AtomRegistry from "effect/reactivity/AtomRegistry";
 import * as Reactivity from "effect/reactivity/Reactivity";
 import { queryCoalesceWindow } from "#queryCoalesceWindow.ts";
 import { liveBatches, liveKeys } from "./liveInvalidation.ts";
@@ -9,7 +10,12 @@ export function queryReactivity<R, E>(runtime: Atom.AtomRuntime<R, E>) {
 	const service = runtime.atom(Effect.service(Reactivity.Reactivity));
 	type Query = Atom.Atom<AsyncResult.AsyncResult<unknown, unknown>>;
 	const cached = new WeakMap<Query, Query>();
-	return (atom: Query, keys: readonly string[]): Query => {
+	return (
+		atom: Query,
+		keys: readonly string[],
+		track: (registry: AtomRegistry.AtomRegistry) => () => void,
+		owner?: AtomRegistry.AtomRegistry,
+	): Query => {
 		const existing = cached.get(atom);
 		if (existing !== undefined) {
 			return existing;
@@ -21,11 +27,15 @@ export function queryReactivity<R, E>(runtime: Atom.AtomRuntime<R, E>) {
 		let wrapped = Atom.transform(
 			source,
 			(get) => {
+				if (owner !== undefined && owner !== get.registry) {
+					throw new Error("A session-scoped query belongs to its binding registry");
+				}
 				const context = get(service);
 				if (!AsyncResult.isSuccess(context)) {
 					return get(source);
 				}
 				const reactivity = context.value;
+				get.addFinalizer(track(get.registry));
 				let pending: (() => boolean) | undefined;
 				let timer: ReturnType<typeof setTimeout> | undefined;
 				function clear() {

package/src/keys.ts

@@ -37,14 +37,26 @@ export function collection<const TName extends string, TId extends Schema.Top &
 	};
 }
 
+export interface KeyContext {
+	readonly scope: string;
+}
+
+function contextual(key: string, context?: KeyContext): string {
+	return context === undefined ? key : `${JSON.stringify(["session", context.scope])}:${key}`;
+}
+
 function own(key: Key): string {
 	return key._tag === "List" ? key.collection : `${key.collection}:${key.id}`;
 }
 
-export function readKeys(keys: readonly Key[]): readonly string[] {
-	return [...new Set(keys.map(own))];
+export function readKeys(keys: readonly Key[], context?: KeyContext): readonly string[] {
+	return [...new Set(keys.flatMap((key) => (key._tag === "Item" ? [own(key), `${key.collection}:*`] : [own(key)])))].map((key) =>
+		contextual(key, context),
+	);
 }
 
-export function invalidationKeys(keys: readonly Key[]): readonly string[] {
-	return [...new Set(keys.flatMap((key) => (key._tag === "List" ? [key.collection] : [own(key), key.collection])))];
+export function invalidationKeys(keys: readonly Key[], context?: KeyContext): readonly string[] {
+	return [...new Set(keys.flatMap((key) => (key._tag === "List" ? [key.collection] : [own(key), key.collection])))].map((key) =>
+		contextual(key, context),
+	);
 }

package/src/live.ts

@@ -1,8 +1,9 @@
 import { type Cause, Duration, Effect, Option, Schema, Stream } from "effect";
 import * as Atom from "effect/reactivity/Atom";
+import type * as AtomRegistry from "effect/reactivity/AtomRegistry";
 import * as Reactivity from "effect/reactivity/Reactivity";
 import { invalidateLive } from "#internal/liveInvalidation.ts";
-import { invalidationKeys, Key, readKeys } from "./keys.ts";
+import { invalidationKeys, Key, type KeyContext, readKeys } from "#keys.ts";
 
 export const LiveHint = Schema.Union([
 	Schema.Struct({ _tag: Schema.Literal("Changed"), keys: Schema.Array(Key) }),
@@ -17,8 +18,9 @@ export interface LiveStatus<E> {
 }
 export interface LiveOptions<E, R> {
 	readonly coalesceWindow?: Duration.Input;
+	readonly keyContext?: KeyContext | undefined;
 	readonly resume?: Atom.Atom<unknown>;
-	readonly resyncKeys: readonly Key[] | (() => readonly Key[]);
+	readonly resyncKeys: readonly Key[] | ((registry: AtomRegistry.AtomRegistry) => readonly Key[]);
 	readonly retryDelay?: Duration.Input;
 	readonly stream: Stream.Stream<LiveHint, E, R>;
 }
@@ -51,14 +53,14 @@ export function live<R, E, TRuntimeError>(runtime: Atom.AtomRuntime<R, TRuntimeE
 				});
 				function changed(changedKeys: readonly Key[]) {
 					const window = get.once(coalesceWindow);
-					const keys = invalidationKeys(changedKeys);
+					const keys = invalidationKeys(changedKeys, options.keyContext);
 					invalidateLive(reactivity, keys, () => active && !get.once(paused), window);
 					const delay = Duration.toMillis(Duration.fromInputUnsafe(window));
 					if (delay === 0) {
 						invalidateLive(reactivity, keys, () => active && !get.once(paused));
 						return;
 					}
-					for (const key of invalidationKeys(changedKeys)) {
+					for (const key of invalidationKeys(changedKeys, options.keyContext)) {
 						pending.add(key);
 					}
 					timer ??= setTimeout(() => {
@@ -75,7 +77,8 @@ export function live<R, E, TRuntimeError>(runtime: Atom.AtomRuntime<R, TRuntimeE
 						update({ needsResync: true });
 						return;
 					}
-					reactivity.invalidateUnsafe(readKeys(typeof options.resyncKeys === "function" ? options.resyncKeys() : options.resyncKeys));
+					const keys = typeof options.resyncKeys === "function" ? options.resyncKeys(get.registry) : options.resyncKeys;
+					reactivity.invalidateUnsafe(keys.map((key) => readKeys([key], options.keyContext)[0] ?? ""));
 					update({ needsResync: false, pending: 0 });
 				}
 				// A subscription alone does not build the flag, so its first read would otherwise announce a change.

package/src/makeLiveServer.ts

@@ -1,13 +1,13 @@
 import { Effect, Queue, Stream } from "effect";
 import * as Reactivity from "effect/reactivity/Reactivity";
 import { makeHintPolicy } from "#internal/makeHintPolicy.ts";
-import { type Key, readKeys } from "./keys.ts";
-import type { LiveHint } from "./live.ts";
+import { type Key, type KeyContext, readKeys } from "#keys.ts";
+import type { LiveHint } from "#live.ts";
 
 export const makeLiveServer = Effect.gen(function* () {
 	const reactivity = yield* Reactivity.Reactivity;
 	const resyncKey = `@shivaedev/effect-contract/resync/${crypto.randomUUID()}`;
-	function stream(keys: readonly Key[]) {
+	function stream(keys: readonly Key[], context?: KeyContext) {
 		return Stream.unwrap(
 			Effect.gen(function* () {
 				const queue = yield* Queue.unbounded<LiveHint>();
@@ -33,7 +33,7 @@ export const makeLiveServer = Effect.gen(function* () {
 					});
 				}
 				const cancel = keys.map((key) => {
-					const own = readKeys([key]);
+					const own = readKeys([key], context);
 					return reactivity.registerUnsafe(own, () => {
 						pending.set(own[0] ?? "", key);
 						flush();

package/src/operation.ts

@@ -1,5 +1,5 @@
 import { Schema } from "effect";
-import type { Bivariant } from "@shivaedev/types/bivariant.ts";
+import type { Bivariant } from "@shivaedev/types/Bivariant.ts";
 import type { Key } from "./keys.ts";
 import { type MatchingTags, type Reject, type RejectionSpecs, type Rejections, type RejectionUnion, rejectionSet } from "./rejection.ts";
 

package/src/serverMetadata.ts

@@ -1,8 +1,8 @@
 import { Context, Schema } from "effect";
 import { Rpc } from "effect/rpc";
-import { Key } from "./keys.ts";
-import { LiveHint } from "./live.ts";
-import type { CommandShape, QueryShape } from "./operation.ts";
+import type { Key } from "#keys.ts";
+import { LiveHint } from "#live.ts";
+import type { CommandShape, QueryShape } from "#operation.ts";
 
 export class QueryReads extends Context.Service<QueryReads, QueryShape["reads"]>()("@shivaedev/effect-contract/QueryReads") {}
 
@@ -10,6 +10,41 @@ export class CommandInvalidation extends Context.Service<CommandInvalidation, Co
 	"@shivaedev/effect-contract/CommandInvalidation",
 ) {}
 
-export function changesRpc<const TName extends string>(name: TName) {
-	return Rpc.make(`${name}.$changes`, { payload: { keys: Schema.Array(Key) }, stream: true, success: LiveHint });
+export class KeyScope extends Context.Service<KeyScope, "session">()("@shivaedev/effect-contract/KeyScope") {}
+
+export interface QueryReference {
+	readonly payload: unknown;
+	readonly query: string;
+}
+
+type ReferenceSchema<TQuery extends QueryShape> = TQuery extends QueryShape
+	? Schema.Struct<{ query: Schema.Literal<TQuery["name"]>; payload: TQuery["payload"] }>
+	: never;
+
+export class ChangesReads extends Context.Service<ChangesReads, (references: readonly QueryReference[]) => readonly Key[]>()(
+	"@shivaedev/effect-contract/ChangesReads",
+) {}
+
+export function queryKeys(queries: readonly QueryShape[], requested: readonly QueryReference[]): readonly Key[] {
+	return requested.flatMap((reference) => {
+		const declared = queries.find((candidate) => candidate.name === reference.query);
+		if (declared === undefined) {
+			throw new Error(`Unknown query: ${reference.query}`);
+		}
+		return declared.reads(reference.payload);
+	});
+}
+
+function referenceSchemas<const TQueries extends readonly QueryShape[]>(
+	queries: TQueries,
+): Schema.$Array<Schema.Union<readonly ReferenceSchema<TQueries[number]>[]>>;
+function referenceSchemas(queries: readonly QueryShape[]) {
+	return Schema.Array(Schema.Union(queries.map((query) => Schema.Struct({ payload: query.payload, query: Schema.Literal(query.name) }))));
+}
+
+export function changesRpc<const TName extends string, const TQueries extends readonly QueryShape[]>(name: TName, queries: TQueries) {
+	return Rpc.make(`${name}.$changes`, { payload: { queries: referenceSchemas(queries) }, stream: true, success: LiveHint }).annotate(
+		ChangesReads,
+		(requested) => queryKeys(queries, requested),
+	);
 }