Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
JavaScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488import assert from "node:assert/strict";import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";import { createServer } from "node:net";import { tmpdir } from "node:os";import { join, resolve } from "node:path";import { test } from "node:test";import { AgentClient } from "agents/client";import { WebSocketChatTransport } from "agents/chat/transport";import { MessageType } from "agents/chat";import WebSocket from "ws";import { unstable_dev } from "wrangler";import { Secret } from "../configuration/secrets.ts";import { createOwnerSession } from "../worker/session.ts";import { customerBindings, installation } from "./fixtures/config.mjs";import { TOOL_ACTIVITY_LIVE_LIMIT } from "../shared/tool-activity.ts";
async function waitFor(predicate) { const deadline = Date.now() + 20_000; while (!(await predicate())) { if (Date.now() > deadline) throw new Error("Timed out waiting for tool activity"); await new Promise((resolve) => setTimeout(resolve, 25)); }}async function freePort() { const server = createServer(); await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); const { port } = server.address(); await new Promise((resolve) => server.close(resolve)); return port;}
test( "native tools publish safe durable activity through execution, cancellation and recovery", { timeout: 180_000 }, async () => { const port = await freePort(); const origin = `http://127.0.0.1:${port}`; const cookie = ( await createOwnerSession( new Secret(customerBindings.FLAREBOT_SESSION_SECRET), { ...installation, runtimeOrigin: origin }, ) ).split(";")[0]; const headers = { Cookie: cookie, Origin: origin }; const persistence = await mkdtemp(join(tmpdir(), "flarebot-activity-")); const config = JSON.parse( await readFile("dist/release/deployment.json", "utf8"), ); const configPath = join(persistence, "wrangler.json"); await writeFile( configPath, JSON.stringify({ ...config, name: "flarebot-activity-test", no_bundle: false, keep_names: true, main: resolve("tests/fixtures/think-worker.ts"), assets: { ...config.assets, directory: resolve("dist/release/assets") }, }), ); const start = () => unstable_dev("tests/fixtures/think-worker.ts", { config: configPath, vars: { ...customerBindings, FLAREBOT_ENV: "development", FLAREBOT_DEV_OVERRIDES: JSON.stringify({ runtimeOrigin: origin }), }, local: true, ip: "127.0.0.1", port, inspectorPort: 0, persist: true, persistTo: persistence, logLevel: "error", experimental: { disableExperimentalWarning: true, watch: false }, }); const clients = []; let worker; try { worker = await start(); class OwnerSocket extends WebSocket { constructor(url, protocols) { super(url, protocols, { headers, closeTimeout: 100 }); } } async function connect(id) { const states = []; const client = new AgentClient({ host: `127.0.0.1:${port}`, protocol: "ws", agent: "PersonalAgent", ...(id ? { basePath: `agents/personal-agent/personal/sub/conversation/${id}`, } : { name: "personal" }), WebSocket: OwnerSocket, onStateUpdate: (state) => states.push(state), }); clients.push(client); const transport = new WebSocketChatTransport({ agent: client }); client.addEventListener("message", (event) => { const frame = JSON.parse(event.data); if (frame.type === MessageType.CF_AGENT_STREAM_RESUMING) transport.handleStreamResuming(frame); if (frame.type === MessageType.CF_AGENT_STREAM_RESUME_NONE) transport.handleStreamResumeNone(frame); if (frame.type === MessageType.CF_AGENT_STREAM_PENDING) transport.handleStreamPending(); }); await client.ready; return { client, transport, states }; } const owner = await connect(); const { id } = await owner.client.call("createConversation", [ "Tool activity", ]); const url = `${origin}/agents/personal-agent/personal/sub/conversation/${id}/get-messages`; const history = async () => (await fetch(url, { headers })).json(); let connection = await connect(id); const observer = await connect(id); const list = async (client = connection.client, before) => client.call("listToolActivities", before === undefined ? [] : [before]); async function send(text) { const chunks = []; const stream = await connection.transport.sendMessages({ chatId: id, trigger: "submit-message", messages: [ ...(await history()), { id: crypto.randomUUID(), role: "user", parts: [{ type: "text", text }], }, ], abortSignal: new AbortController().signal, }); const done = (async () => { for await (const chunk of stream) chunks.push(chunk); })(); return { done, chunks }; } assert.deepEqual(await list(), { activities: [], nextCursor: null }); for (const invalid of [null, -1, 0, "1", 1.5, {}]) await assert.rejects( connection.client.call("listToolActivities", [invalid]), /Invalid tool activity cursor/, ); await assert.rejects( connection.client.call("reportToolProgress", ["forged", {}]), /not callable/, ); await assert.rejects( connection.client.call("setState", [{ toolActivities: [] }]), /not callable/, ); assert.equal((await fetch(url)).status, 401); assert.equal( ( await fetch(url, { headers: { ...headers, Origin: "https://attacker.invalid" }, }) ).status, 403, );
const pending = await send("activity-pending"); await waitFor(() => observer.states.some((s) => s.toolActivities.some((a) => a.status === "pending"), ), ); await pending.done; const first = (await list()).activities[0]; assert.equal(first.status, "succeeded"); assert.equal(first.kind, "web"); assert.equal(first.inputSummary, "Research requested"); assert.ok( first.createdAt <= first.startedAt && first.startedAt <= first.endedAt, ); assert.ok(first.endedAt === first.updatedAt); const statesFor = (callId) => observer.states.flatMap((s) => s.toolActivities.filter((a) => a.toolCallId === callId), ); assert.ok( statesFor(first.toolCallId).some((a) => a.status === "running"), ); connection.client.setState({ toolActivityVersion: 1, toolActivities: [{ toolCallId: "forged-secret" }], }); await new Promise((resolve) => setTimeout(resolve, 100)); assert.deepEqual((await list()).activities, [first]); assert.ok(!JSON.stringify(observer.states).includes("forged-secret"));
for (const scenario of [ "activity-error", "activity-failure", "activity-action", "activity-unknown", ]) { const turn = await send(scenario); await turn.done; const activity = (await list()).activities[0]; assert.equal( activity.status, scenario === "activity-unknown" ? "succeeded" : "failed", scenario, ); if (scenario === "activity-unknown") { assert.equal(activity.toolName, "unknown"); assert.equal(activity.inputSummary, "Tool invocation"); } else assert.equal(activity.reason, "tool-error"); } const boundedTurn = await send("activity-bounds"); await boundedTurn.done; const bounded = (await list()).activities[0]; assert.equal(bounded.inputSummary.length, 160); assert.ok(!bounded.inputSummary.includes("\n")); const early = await send("activity-early"); await waitFor(async () => (await list()).activities[0]?.progress); assert.equal( (await list()).activities[0].status, "running", "a single early progress update publishes before silent work finishes", ); await early.done; const parallel = await send("activity-parallel"); await waitFor(() => observer.states.some( (s) => s.toolActivities.filter((a) => a.status === "running").length === 2, ), ); await waitFor(() => observer.states.some( (s) => s.toolActivities[0]?.status === "running" && s.toolActivities.some((a) => a.status === "succeeded"), ), ); await parallel.done; const parallelResults = (await list()).activities.slice(0, 2); assert.ok(parallelResults.every((a) => a.status === "succeeded")); assert.notEqual( parallelResults[0].toolCallId, parallelResults[1].toolCallId, );
const progress = await send("activity-progress"); await waitFor(() => observer.states.some((s) => s.toolActivities.some((a) => a.status === "running" && a.progress), ), ); await progress.done; const progressed = (await list()).activities[0]; assert.equal(progressed.status, "succeeded"); assert.equal(progressed.progress.text, "Processing items"); assert.ok( statesFor(progressed.toolCallId).filter((a) => a.progress).length < 20, "1000 preliminary results are coalesced", ); assert.ok( progress.chunks.some( (c) => c.type === "tool-output-available" && c.preliminary, ), "native preliminary chunks remain intact", );
const pendingCancel = await send("activity-pending"); pendingCancel.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "pending", ); const pendingCancelId = (await list()).activities[0].toolCallId; connection.transport.cancelActiveServerTurn(); await assert.rejects(pendingCancel.done, { name: "AbortError" }); await waitFor( async () => (await list()).activities.find( (a) => a.toolCallId === pendingCancelId, )?.status === "cancelled", );
for (const scenario of ["activity-cancel", "activity-late"]) { const turn = await send(scenario); turn.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "running", ); const callId = (await list()).activities[0].toolCallId; assert.equal(connection.transport.cancelActiveServerTurn(), true); await assert.rejects(turn.done, { name: "AbortError" }); await waitFor( async () => (await list()).activities.find((a) => a.toolCallId === callId) ?.status === "cancelled", ); const next = await send("tool"); await next.done; await new Promise((resolve) => setTimeout(resolve, 1700)); const cancelled = (await list()).activities.find( (a) => a.toolCallId === callId, ); assert.equal(cancelled.status, "cancelled"); assert.equal( cancelled.outputSummary, "Conversation stopped waiting for this tool", ); const observed = statesFor(callId).map((a) => a.status); assert.ok( observed .slice(observed.indexOf("cancelled")) .every((s) => s === "cancelled"), "late completion cannot regress terminal state", ); } const detached = await send("activity-detach"); detached.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "running", ); const detachedId = (await list()).activities[0].toolCallId; connection.client.close(); connection = await connect(id); await waitFor(() => connection.states.length); assert.equal( connection.states .at(-1) .toolActivities.find((a) => a.toolCallId === detachedId)?.status, "running", ); const resumed = await connection.transport.reconnectToStream({ chatId: id, }); if (resumed) for await (const _ of resumed) { /* native drain */ } await waitFor( async () => (await list()).activities.find((a) => a.toolCallId === detachedId) ?.status === "succeeded", );
const many = await send("activity-many"); await many.done; await waitFor(() => observer.states .at(-1) .toolActivities.every((a) => a.status === "succeeded"), ); const recent = await list(); assert.equal(recent.activities.length, 100); assert.ok(recent.nextCursor); const older = await list(connection.client, recent.nextCursor); assert.deepEqual( older.activities.find((a) => a.toolCallId === first.toolCallId), first, "older timestamps and summaries survive live-window eviction", ); const all = [...recent.activities, ...older.activities]; assert.equal(new Set(all.map((a) => a.toolCallId)).size, all.length); assert.ok( observer.states.every( (s) => s.toolActivities.length <= TOOL_ACTIVITY_LIVE_LIMIT, ), ); for (const secret of [ "fixture-private", "user:password", "token=secret", "private?", customerBindings.FLAREBOT_SESSION_SECRET, ]) { assert.ok(!JSON.stringify(all).includes(secret)); assert.ok(!JSON.stringify(observer.states).includes(secret)); } const recovery = await send("activity-recover"); recovery.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "running", ); const interruptedId = (await list()).activities[0].toolCallId; await new Promise((resolve) => setTimeout(resolve, 400)); clients.forEach((client) => client.close()); await worker.stop(); worker = await start(); connection = await connect(id); const interrupted = (await list()).activities.find( (a) => a.toolCallId === interruptedId, ); assert.equal(interrupted.status, "failed"); assert.equal(interrupted.reason, "interrupted"); assert.equal( interrupted.outputSummary, "Execution interrupted; outcome unknown", ); await waitFor(async () => JSON.stringify(await history()).includes("activity-recover complete"), ); const restartedPage = await list(); const restartedOlder = await list( connection.client, restartedPage.nextCursor, ); assert.deepEqual( restartedOlder.activities.find( (a) => a.toolCallId === first.toolCallId, ), first, ); assert.ok( [...restartedPage.activities, ...restartedOlder.activities].every( (a) => !["running", "pending"].includes(a.status), ), ); const replay = await send("activity-replay"); replay.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "running", ); const replayBefore = (await list()).activities[0]; await new Promise((resolve) => setTimeout(resolve, 400)); clients.forEach((client) => client.close()); await worker.stop(); worker = await start(); connection = await connect(id); await waitFor( async () => (await list()).activities.find( (a) => a.toolCallId === replayBefore.toolCallId, )?.status === "succeeded", ); const replayAfter = (await list()).activities.find( (a) => a.toolCallId === replayBefore.toolCallId, ); assert.equal(replayAfter.attempts, 2); assert.equal(replayAfter.createdAt, replayBefore.createdAt); assert.ok(replayAfter.interruptedAt); assert.equal(replayAfter.reason, undefined); assert.ok( (await history()) .flatMap((m) => m.parts) .some( (p) => p.toolCallId === replayAfter.toolCallId && p.state === "output-available", ), "recovered native transcript and same-ID activity agree", );
const clearTurn = await send("activity-late"); clearTurn.done.catch(() => {}); await waitFor( async () => (await list()).activities[0]?.status === "running", ); connection.client.send( JSON.stringify({ type: MessageType.CF_AGENT_CHAT_CLEAR }), ); await waitFor(async () => (await list()).activities.length === 0); await waitFor(async () => (await history()).length === 0); const afterClear = await send("tool"); await afterClear.done; await new Promise((resolve) => setTimeout(resolve, 1700)); assert.equal( (await list()).activities.length, 1, "late pre-clear tool completion cannot recreate any summary or erase the new turn", ); assert.equal((await list()).activities[0].toolName, "fixtureEcho");
const finalTurn = await send("tool"); await finalTurn.done; assert.equal((await list()).activities[0].status, "succeeded"); } finally { clients.forEach((client) => client.close()); await worker?.stop(); await rm(persistence, { recursive: true, force: true }); } },);