Skip to content
Merged
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
8 changes: 8 additions & 0 deletions .changeset/coupled-child-proofs.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
---
"@statelyai/agent": patch
---

`children: "coupled"` exploration closes two proof gaps:

- A `canReach` predicate and `dedupeBy` now see each coupled child's explored state through `snapshot.children[id].getSnapshot()`, not the state it started in. A parent's send and the child's receipt are one state, so cross-actor invariants can be targets.
- A coupled child's text requests deliver `@agent.usage` to the root, as a run does: each `usage` value forks a branch, and the root takes or drops the event by the usual rules. Token budgets fed by child calls can now be proved.
5 changes: 4 additions & 1 deletion docs/verify.md
Original file line number Diff line number Diff line change
Expand Up @@ -313,7 +313,8 @@ const result = await canReach(waitingMachine, "timedOut");
- A usage value the active states refuse (`snapshot.can` is false) is dropped, as a run drops it: the request settles with no usage event delivered or recorded.
- A usage transition that leaves the invoking state, or restarts its invoke, cancels the request: its completion is not recorded, and a restarted invoke settles later on its own.
- A settling request with no `usage` entry while a usage transition is active lists `@agent.usage` in `vacuous`. Completeness is unchanged.
- Requests inside child machines and decisions deliver no usage.
- With `children: "coupled"`, a child's text request with an entry forks one branch per value. The usage event goes to the root before the child's completion reaches its parent. The root takes it or drops it by the same rules. A child's request with no entry is not listed in `vacuous`.
- Requests inside `"outputs"` or `"explore"` children, and decisions, deliver no usage.

```ts no-check
const overBudget = await canReach(researchMachine, "outOfBudget", {
Expand Down Expand Up @@ -361,7 +362,9 @@ const report = await explorePaths(parentMachine, {
- The parent's host events arrive at its idle waits, including while every child waits on the host. While a child or the parent has work pending, they race it only with `interleave: "pending"`.
- What a child's move sends its own children, and their replies, happen in the same move. A child's send to a sibling is delivered to that sibling, entry sends included. A child's sends reach their targets in the order it made them, each run to quiescence before the next, as a live run does. A send to an actor the search does not model lists the event type in `unexpandedEvents`, so the report is not complete.
- Child moves count against `maxDepth` like the parent's own branches. On a path they carry `actor` with the child's id (a grandchild's as `child.grandchild`): the host events it received, its choices, its settled work. The events it sends the parent appear as the parent receives them.
- A target predicate and `dedupeBy` see each coupled child's explored state: `snapshot.children[id].getSnapshot()` returns the state the search stepped the child to, its own coupled children likewise. A parent's send and the child's receipt are one state, so a cross-actor invariant such as "every chunk the parent forwarded, a child received" can be a target.
- With `dedupeBy`, each child's state joins the key. Without it, every interleaving is explored in full.
- A coupled child's text requests deliver `@agent.usage` to the root, as a run does. See [Usage events](#usage-events).
- A child that stops on `needs-output` or `max-depth` ends the parent branch the same way, so the parent report is not complete.

Pass `hostEvents` when some types the parent accepts come only from its children, such as a child's status events. The search then offers only the listed types at the parent's waits.
Expand Down
227 changes: 225 additions & 2 deletions src/verify-coupled.test.ts
Original file line number Diff line number Diff line change
@@ -1,8 +1,8 @@
import { describe, expect, test } from "vitest";
import { z } from "zod";
import { createAsyncLogic, createMachine, type AnyStateMachine } from "xstate";
import { setupAgent } from "./index.js";
import { canReach, explorePaths } from "./testing/index.js";
import { AGENT_USAGE_EVENT_TYPE, runAgent, setupAgent } from "./index.js";
import { canReach, explorePaths, type ExplorePathsOptions } from "./testing/index.js";
import { isExplorationComplete } from "./verify.js";

type Enq = { sendTo: (to: unknown, event: unknown) => void };
Expand Down Expand Up @@ -745,3 +745,226 @@ describe("explorePaths — children: 'coupled'", () => {
});
});
});

describe("coupled child snapshots in predicates", () => {
const counter = setupAgent({
context: z.object({ n: z.number() }),
events: { PING: z.object({}) },
}).createMachine({
context: { n: 0 },
initial: "a",
states: {
a: { on: { PING: ({ context }) => ({ target: "b", context: { n: context.n + 1 } }) } },
b: {},
},
});

const pinger = setupAgent({
context: z.object({ sent: z.number() }),
events: { GO: z.object({}) },
actors: { counter },
}).createMachine({
context: { sent: 0 },
entry: ({ actors }, enq) => {
enq.spawn(actors.counter, { id: "c" });
},
on: {
GO: ({ context, children }, enq) => {
if (context.sent > 0) return undefined;
enq.sendTo(children.c, { type: "PING" });
return { context: { sent: context.sent + 1 } };
},
},
});

const childState = (snapshot: any, id: string) => snapshot.children[id]?.getSnapshot();

test("a coupled predicate sees the child's explored state", async () => {
const seen = new Set<unknown>();
const reached = await canReach(
pinger,
(snapshot) => {
seen.add(childState(snapshot, "c")?.value);
return childState(snapshot, "c")?.value === "b";
},
{ children: "coupled", maxDepth: 3 },
);
expect(reached.reachable).toBe(true);
expect(reached.witness).toEqual([{ type: "GO" }]);
expect(seen).toEqual(new Set(["a", "b"]));
});

test("a parent send and the child's receipt are one state to a predicate", async () => {
// sent === child n on every explored state: the send applies at once.
const broken = await canReach(
pinger,
(snapshot) => snapshot.context.sent !== childState(snapshot, "c")?.context.n,
{ children: "coupled", maxDepth: 3 },
);
expect(broken.reachable).toBe(false);
});

const stepper = setupAgent({ events: { GO: z.object({}) }, context: z.object({}) }).createMachine(
{
context: {},
initial: "a",
states: {
a: { on: { GO: { target: "b" } } },
b: { on: { GO: { target: "c" } } },
c: { type: "final" },
},
},
);

const host = setupAgent({ events: {}, context: z.object({}), actors: { stepper } }).createMachine(
{
context: {},
initial: "running",
states: {
running: { invoke: { id: "kid", src: "stepper", onDone: { target: "done" } } },
done: { type: "final" },
},
},
);

test("a predicate sees an invoked child's own moves", async () => {
const reached = await canReach(host, (snapshot) => childState(snapshot, "kid")?.value === "b", {
children: "coupled",
childEvents: { stepper: ["GO"] },
});
expect(reached.reachable).toBe(true);
expect(reached.witness).toEqual([{ type: "GO", actor: "kid" }]);
});

test("dedupeBy sees the coupled child's explored state", async () => {
const keys = new Set<string>();
await explorePaths(host, {
children: "coupled",
childEvents: { stepper: ["GO"] },
dedupeBy: (snapshot) => {
const key = `${JSON.stringify(snapshot.value)}:${String(childState(snapshot, "kid")?.value)}`;
keys.add(key);
return key;
},
});
expect(keys).toEqual(new Set(['"running":a', '"running":b', '"done":undefined']));
});
});

describe("coupled child usage", () => {
const workerAgent = setupAgent({
context: z.object({}),
events: {},
requests: {
summarize: { model: "m", schemas: { input: z.object({}), output: z.string() }, prompt: "x" },
},
});
const worker = workerAgent.createMachine({
context: {},
initial: "working",
states: {
working: { invoke: { id: "s", src: "summarize", input: {}, onDone: { target: "done" } } },
done: { type: "final" },
},
});
// A token budget the child's request feeds: the root is charged, as in a run.
const createBudget = (onUsage?: (event: unknown) => void) =>
setupAgent({
context: z.object({ tokens: z.number() }),
events: {},
actors: { worker },
}).createMachine({
context: { tokens: 0 },
initial: "running",
states: {
running: {
on: {
[AGENT_USAGE_EVENT_TYPE]: ({ context, event }) => {
onUsage?.(event);
const tokens = context.tokens + (event.usage.totalTokens ?? 0);
return tokens > 100
? { target: "overBudget", context: { tokens } }
: { context: { tokens } };
},
},
invoke: { id: "w", src: "worker", onDone: { target: "done" } },
},
overBudget: { type: "final" },
done: { type: "final" },
},
});
const usageEvent = (totalTokens: number) => ({
type: AGENT_USAGE_EVENT_TYPE,
kind: "text",
id: "s",
src: "summarize",
model: "m",
name: "summarize",
usage: { totalTokens },
});
const options: ExplorePathsOptions = {
children: "coupled",
outputs: { summarize: ["ok"] },
usage: { summarize: [{ totalTokens: 10 }, { totalTokens: 500 }] },
};

test("a child request's usage is delivered to the root before its completion", async () => {
const report = await explorePaths(createBudget(), options);
expect(report.complete).toBe(true);
expect(report.terminals).toEqual([
{
status: "done",
path: [
{ type: "xstate.done.actor", actorId: "s", output: "ok", actor: "w" },
usageEvent(10),
{ type: "xstate.done.actor", actorId: "w" },
],
state: "done",
},
{
status: "done",
path: [
{ type: "xstate.done.actor", actorId: "s", output: "ok", actor: "w" },
usageEvent(500),
],
state: "overBudget",
},
]);
expect(
await canReach(createBudget(), (snapshot) => snapshot.context.tokens > 100, options),
).toMatchObject({ reachable: true });
});

test("the root drops usage it does not take", async () => {
const plain = setupAgent({
context: z.object({}),
events: {},
actors: { worker },
}).createMachine({
context: {},
initial: "running",
states: {
running: { invoke: { id: "w", src: "worker", onDone: { target: "done" } } },
done: { type: "final" },
},
});
const report = await explorePaths(plain, options);
expect(report.complete).toBe(true);
expect(report.vacuous).toEqual([]);
for (const terminal of report.terminals) {
expect(terminal.path.map((event) => event.type)).not.toContain(AGENT_USAGE_EVENT_TYPE);
expect(terminal.state).toBe("done");
}
});

test("a run charges the root with the same event, bar its budget and call key", async () => {
const seen: unknown[] = [];
const run = await runAgent(
createBudget((event) => seen.push(event)),
{ executors: { generateText: async () => ({ result: "ok", usage: { totalTokens: 10 } }) } },
);
expect(run.status).toBe("done");
const { budget: _budget, callKey: _callKey, ...live } = seen[0] as Record<string, unknown>;
expect(live).toEqual(usageEvent(10));
});
});
Loading
Loading