refactor(agent-kit): pass CAS store through AgentContext

Expose the store created during context build on AgentContext so agents
reuse the same in-memory cache instead of opening a second store.

Co-authored-by: Cursor <cursoragent@cursor.com>

packages/uwf-agent-hermes/src/hermes.ts

@@ -1,12 +1,6 @@
 import { spawn } from "node:child_process";
 
-import {
+import { type AgentContext, type AgentRunResult, createAgent } from "@uncaged/uwf-agent-kit";
-  type AgentContext,
-  type AgentRunResult,
-  createAgent,
-  createAgentStore,
-  resolveStorageRoot,
-} from "@uncaged/uwf-agent-kit";
 
 import {
   loadHermesSession,
@@ -92,8 +86,7 @@ function spawnHermesChat(prompt: string): Promise<string> {
 async function runHermes(ctx: AgentContext): Promise<AgentRunResult> {
   const fullPrompt = buildHermesPrompt(ctx);
   const rawOutput = await spawnHermesChat(fullPrompt);
-  const storageRoot = resolveStorageRoot();
+  const { store } = ctx;
-  const { store } = await createAgentStore(storageRoot);
 
   const sessionId = parseSessionIdFromStdout(rawOutput);
   if (sessionId !== null) {

packages/uwf-agent-kit/src/context.ts

@@ -1,3 +1,4 @@
+import type { Store } from "@uncaged/json-cas";
 import type {
   CasRef,
   StartNodePayload,
@@ -6,6 +7,7 @@ import type {
   ThreadId,
 } from "@uncaged/uwf-protocol";
 import { createAgentStore, loadThreadsIndex, resolveStorageRoot } from "./storage.js";
+import type { AgentStore } from "./storage.js";
 import type { AgentContext } from "./types.js";
 
 type ChainState = {
@@ -20,8 +22,8 @@ function fail(message: string): never {
 }
 
 function walkChain(
-  store: Awaited<ReturnType<typeof createAgentStore>>["store"],
+  store: Store,
-  schemas: Awaited<ReturnType<typeof createAgentStore>>["schemas"],
+  schemas: AgentStore["schemas"],
   headHash: CasRef,
 ): ChainState {
   const headNode = store.get(headHash);
@@ -77,7 +79,7 @@ function walkChain(
 }
 
 function expandOutput(
-  store: Awaited<ReturnType<typeof createAgentStore>>["store"],
+  store: Store,
   outputRef: CasRef,
 ): unknown {
   const node = store.get(outputRef);
@@ -88,7 +90,7 @@ function expandOutput(
 }
 
 async function buildHistory(
-  store: Awaited<ReturnType<typeof createAgentStore>>["store"],
+  store: Store,
   stepsNewestFirst: StepNodePayload[],
 ): Promise<StepContext[]> {
   const chronological = [...stepsNewestFirst].reverse();
@@ -105,8 +107,8 @@ async function buildHistory(
 }
 
 async function loadWorkflow(
-  store: Awaited<ReturnType<typeof createAgentStore>>["store"],
+  store: Store,
-  schemas: Awaited<ReturnType<typeof createAgentStore>>["schemas"],
+  schemas: AgentStore["schemas"],
   workflowRef: CasRef,
 ) {
   const node = store.get(workflowRef);
@@ -150,13 +152,14 @@ export async function buildContext(threadId: ThreadId, role: string): Promise<Ag
     prompt: chain.start.prompt,
     history,
     workflow,
+    store,
   };
 }
 
 export type BuildContextMeta = {
   storageRoot: string;
-  store: Awaited<ReturnType<typeof createAgentStore>>["store"];
+  store: Store;
-  schemas: Awaited<ReturnType<typeof createAgentStore>>["schemas"];
+  schemas: AgentStore["schemas"];
   headHash: CasRef;
   chain: ChainState;
 };
@@ -194,6 +197,7 @@ export async function buildContextWithMeta(
     prompt: chain.start.prompt,
     history,
     workflow,
+    store,
     meta: { storageRoot, store, schemas, headHash, chain },
   };
 }

packages/uwf-agent-kit/src/index.ts

@@ -7,11 +7,5 @@ export {
   resolveModel,
 } from "./extract.js";
 export { createAgent } from "./run.js";
-export {
+export { getConfigPath, getEnvPath, loadWorkflowConfig } from "./storage.js";
-  createAgentStore,
-  getConfigPath,
-  getEnvPath,
-  loadWorkflowConfig,
-  resolveStorageRoot,
-} from "./storage.js";
 export type { AgentContext, AgentOptions, AgentRunFn, AgentRunResult } from "./types.js";

packages/uwf-agent-kit/src/types.ts

@@ -1,3 +1,4 @@
+import type { Store } from "@uncaged/json-cas";
 import type { StepContext, ThreadId, WorkflowPayload } from "@uncaged/uwf-protocol";
 
 export type AgentContext = {
@@ -7,6 +8,7 @@ export type AgentContext = {
   prompt: string;
   history: StepContext[];
   workflow: WorkflowPayload;
+  store: Store;
 };
 
 export type AgentRunResult = {

scripts/mock-agent.ts

@@ -1,11 +1,7 @@
 #!/usr/bin/env bun
 // Mock agent for smoke testing
 import { bootstrap, type JSONSchema, putSchema } from "@uncaged/json-cas";
-import {
+import { createAgent } from "../packages/uwf-agent-kit/src/index.js";
-  createAgent,
-  createAgentStore,
-  resolveStorageRoot,
-} from "../packages/uwf-agent-kit/src/index.js";
 
 const MOCK_RAW_OUTPUT_SCHEMA: JSONSchema = {
   title: "mock-raw-output",
@@ -21,7 +17,7 @@ const agent = createAgent({
   name: "mock",
   run: async (ctx) => {
     const output = `Mock output for role ${ctx.role}: task was "${ctx.prompt}"`;
-    const { store } = await createAgentStore(resolveStorageRoot());
+    const { store } = ctx;
     await bootstrap(store);
     const schemaHash = await putSchema(store, MOCK_RAW_OUTPUT_SCHEMA);
     const detailHash = await store.put(schemaHash, { text: output });