Skip to content

Load agent catalogs each turn instead of caching them on the chat #267

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
AshishKumar4 wants to merge 1 commit into
base: main
Choose a base branch
from
+17 −99
Open

Load agent catalogs each turn instead of caching them on the chat #267

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
40 changes: 1 addition & 39 deletions packages/workshop-backend/__tests__/agent-catalog.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,7 @@ import {
AGENT_CATALOG_MAX_TITLE_LENGTH, boundAgentCatalog,
} from "@gadgets/workshop-shared/gatekeeper";
import {
completeAgentCatalogSnapshot, formatAgentCatalogPrompt,
formatAlwaysAvailableResourcesPrompt, normalizeAgentCatalog,
formatAgentCatalogPrompt, formatAlwaysAvailableResourcesPrompt, normalizeAgentCatalog,
} from "../src/agent-catalog";

describe("normalizeAgentCatalog", () => {
Expand Down Expand Up @@ -129,42 +128,5 @@ describe("boundAgentCatalog", () => {
});
});

describe("completeAgentCatalogSnapshot", () => {
it("loads each catalog once and preserves null snapshots", async () => {
let calls: number[] = [];
let first = await completeAgentCatalogSnapshot(undefined, [2, 1], async gatekeeperId => {
calls.push(gatekeeperId);
return gatekeeperId === 1 ? {entries: []} : null;
});
let second = await completeAgentCatalogSnapshot(first.snapshots, [2, 1], async gatekeeperId => {
calls.push(gatekeeperId);
return {entries: []};
});

expect(first).toEqual({
snapshots: [
{gatekeeperId: 1, catalog: {entries: []}},
{gatekeeperId: 2, catalog: null},
],
changed: true,
});
expect(second).toEqual({snapshots: first.snapshots, changed: false});
expect(calls.toSorted()).toEqual([1, 2]);
});

it("prunes snapshots for removed gatekeepers", async () => {
let result = await completeAgentCatalogSnapshot([
{gatekeeperId: 1, catalog: {entries: []}},
{gatekeeperId: 2, catalog: null},
], [1], async () => {
throw new Error("existing catalogs must not be reloaded");
});

expect(result).toEqual({
snapshots: [{gatekeeperId: 1, catalog: {entries: []}}],
changed: true,
});
});
});


38 changes: 0 additions & 38 deletions packages/workshop-backend/src/agent-catalog.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,6 @@ import { createWorkshopLogger } from "./observability";

const logger = createWorkshopLogger("workshop.agent.catalog");

export type AgentCatalogSnapshot = {
gatekeeperId: number;
catalog: AgentCatalog | null;
};

function normalizeText(value: string, maxLength: number): string {
return value.replace(/\p{Cc}/gu, " ").replace(/\s+/g, " ").trim().slice(0, maxLength);
}
Expand Down Expand Up @@ -46,39 +41,6 @@ export function normalizeAgentCatalog(catalog: AgentCatalog): AgentCatalog {
};
}

export async function completeAgentCatalogSnapshot(
existing: AgentCatalogSnapshot[] | undefined,
gatekeeperIds: number[],
loadCatalog: (gatekeeperId: number) => Promise<AgentCatalog | null>):
Promise<{snapshots: AgentCatalogSnapshot[], changed: boolean}> {
let activeIds = new Set(gatekeeperIds);
let existingCount = existing?.length ?? 0;
let catalogs = new Map(
existing
?.filter(entry => activeIds.has(entry.gatekeeperId))
.map(entry => [entry.gatekeeperId, entry.catalog]));
let removedStaleEntries = catalogs.size !== existingCount;
let missing = gatekeeperIds.filter(gatekeeperId => !catalogs.has(gatekeeperId));
await Promise.all(missing.map(async gatekeeperId => {
// Isolate per entry: one failing loader must not reject the whole snapshot (it would lose every
// other catalog and abort the turn). A failed/empty load is recorded as null, like any other.
try {
catalogs.set(gatekeeperId, await loadCatalog(gatekeeperId));
} catch (error) {
logger.warn("failed to load agent catalog", {
event: "agent.catalog.load.failed", gatekeeperId, error,
});
catalogs.set(gatekeeperId, null);
}
}));
return {
snapshots: [...catalogs]
.toSorted(([left], [right]) => left - right)
.map(([gatekeeperId, catalog]) => ({gatekeeperId, catalog})),
changed: missing.length > 0 || removedStaleEntries,
};
}

/** The catalog as a JSON blob for inclusion in a prompt, on its own line, or "" if empty. */
export function formatAgentCatalogPrompt(catalog: AgentCatalog | null): string {
if (!catalog?.entries.length) return "";
Expand Down
7 changes: 1 addition & 6 deletions packages/workshop-backend/src/agent.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@ import {
import { RpcStub as NativeRpcStub } from "cloudflare:workers";
import { createTwoFilesPatch, FILE_HEADERS_ONLY } from "diff";
import { webFetch as webFetchImpl, WebFetchEnv, formatWebFetchResult } from "./web-fetch";
import { AgentCatalogSnapshot, formatAlwaysAvailableResourcesPrompt } from "./agent-catalog";
import { formatAlwaysAvailableResourcesPrompt } from "./agent-catalog";
import { formatInstanceInstructions } from "./admin-config";
import type { AiGatewayLogRoute } from "./ai-gateway";
import { AgentTurnError, completeText, httpStatusFromError, zeroUsage } from "./ai-invoke";
Expand Down Expand Up @@ -67,11 +67,6 @@ export type AiChatAgentContext = {
*/
alwaysAvailableCapsuleIds?: WorkpieceId[];

/**
* Cached discovery catalogs for the always-available resources, keyed per gatekeeper.
* Regenerable: re-fetched when missing/stale (see prepareChatBindings).
*/
alwaysAvailableCatalogs?: AgentCatalogSnapshot[];
};

/**
Expand Down
31 changes: 15 additions & 16 deletions packages/workshop-backend/src/overseer.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import { RpcCompatible, RpcStub, RpcTarget } from "capnweb";
import { validateRpc } from "capnweb-validate";
import { Overseer, GadgetMetadata, UiBundle, WorkpieceId, WorkpieceSummary, WorkpiecesSubscriber, GadgetClient, GadgetBindingInfo, GatekeeperClient, ActionState, ActionLogEntry, ActionsSubscriber, CodeUpdate, CodeSubscriber, AiChatMetadata, AiChatMessage, AiChatHistoryPage, AiChatSubscriber, AiChatAuthorInfo, AiModelConfig, AiChatMessageBody, AgentSpawnerConfig, ConsoleLogSubscriber, ConsoleLogEvent, CapsuleSpecifier, CollaboratorInfo, CollaboratorRole, AffectedCollaborator, ShareLinkInfo, GatekeeperCreationSpec, ObserverConfigCallback, ObserverBindingNeed, ObserverBindingFailure, BlueprintBindingAnnotation, BlueprintBinding, BlueprintMetadata, BlueprintOutput, MessageFormatRef, isOutputIcon, SpawnerEnvTarget, BlueprintGadgetSummary, AiChatStreamEvent, BlueprintScreenshotUpload, BLUEPRINT_SCREENSHOT_R2_PREFIX, blueprintScreenshotUrl, ChatAttachmentUpload, ChatAttachmentHandle, ChatAttachmentRef, BoundHookInfo, PreApprovableAction, PresenceParticipant, PresenceSubscriber, SlashCommandChoice, SlashCommandRequest, validateBindingName, createOpenGadgetError, OPEN_GADGET_ERROR_CODES, resolveSiteName } from '@gadgets/workshop-shared/api';
import { Gatekeeper, HookInitiator, ResourceDescription, ApprovalQueue, ActionDescription, ObservationAuthorizer, ObservationDescription, VendorDescription, SupportedResource, resolveRequestedResource, HookController, HookDescription, ActionKind } from "@gadgets/workshop-shared/gatekeeper";
import { AgentCatalog, Gatekeeper, HookInitiator, ResourceDescription, ApprovalQueue, ActionDescription, ObservationAuthorizer, ObservationDescription, VendorDescription, SupportedResource, resolveRequestedResource, HookController, HookDescription, ActionKind } from "@gadgets/workshop-shared/gatekeeper";
import {
DurableObject, WorkerEntrypoint, RpcStub as NativeRpcStub,
RpcTarget as NativeRpcTarget, restore,
Expand Down Expand Up @@ -32,7 +32,7 @@ import { recordAnalytics } from "./analytics";
import { reportIssue } from "@gadgets/backend-utils/error-reporting";
import type { ProductAnalyticsConnectionType, ProductAnalyticsGadgetInput } from "./analytics";
import { checkUsageAndBalance } from "./ai-gateway-billing/limits/usage-checker";
import { completeAgentCatalogSnapshot, normalizeAgentCatalog } from "./agent-catalog";
import { normalizeAgentCatalog } from "./agent-catalog";
import { refreshCachedBalance } from "./ai-gateway-billing/cloudflare/connection-service";
import { SharingManager, SharingCaller, CollaboratorRecord, ShareKeyRecord } from "./sharing";
import { AutoApprovalDrainer } from "./auto-approval";
Expand Down Expand Up @@ -4940,13 +4940,16 @@ class OverseerImpl implements AgentHooks {
}
}

// Complete/refresh the cached discovery catalogs for the frozen ambient set.
let {snapshots, changed} = await completeAgentCatalogSnapshot(
context.alwaysAvailableCatalogs,
ambientIds,
async gatekeeperId => {
// Load the discovery catalogs for the frozen ambient set.
//
// Deliberately not cached on the chat. A catalog says what the session can reach *now*, so a
// cached one can never show a skill added after the chat opened, and a cached failure reads as
// an empty library for the rest of the chat. Rebuilding it costs one call per ambient
// gatekeeper per turn, which is the same call the chat already makes to build its bindings.
let catalogs = new Map(await Promise.all(ambientIds.map(
async (gatekeeperId): Promise<[number, AgentCatalog | null]> => {
let record = this.storage.gatekeepers.get(gatekeeperId);
if (!record) return null; // disconnected since the chat froze its set — no catalog.
if (!record) return [gatekeeperId, null]; // disconnected since the chat froze its set.
try {
using authorizer = new RpcStub<ObservationAuthorizer>(new ApprovalQueueImpl(
this, gatekeeperId, {from: "agent", chatId}));
Expand All @@ -4959,7 +4962,7 @@ class OverseerImpl implements AgentHooks {
let facet = this.getGatekeeperFacet(gatekeeperId) as unknown as CatalogGatekeeperFacet;
let catalog = await facet.getAgentCatalog(
authorizer as unknown as ObservationAuthorizer);
return catalog ? normalizeAgentCatalog(catalog) : null;
return [gatekeeperId, catalog ? normalizeAgentCatalog(catalog) : null];
} catch (error) {
reportIssue("overseer.catalog-fallback", error, {
handled: true,
Expand All @@ -4971,13 +4974,10 @@ class OverseerImpl implements AgentHooks {
event: "agent.catalog.load.failed",
gatekeeperId, resourceTitle: record.resourceTitle, error,
});
return null;
// The next turn loads it again, so one failure costs this turn's catalog and no more.
return [gatekeeperId, null];
}
});
if (changed) {
context.alwaysAvailableCatalogs = snapshots;
dirty = true;
}
})));
if (dirty) {
// The work above is async, so the chat could have been deleted meanwhile. Don't resurrect
// its per-chat storage: deleteChat is the single cleanup point (see its comment) and
Expand All @@ -4989,7 +4989,6 @@ class OverseerImpl implements AgentHooks {

// Materialize the seed entries, skipping targets that no longer exist (mirroring env build);
// ambient entries carry their catalogs.
let catalogs = new Map(snapshots.map(entry => [entry.gatekeeperId, entry.catalog]));
let ambientSet = new Set(ambientIds);
let result: SeedBindingInfo[] = [];
for (let [name, target] of Object.entries(seedMap)) {
Expand Down
Loading