@poncho-ai/harness 0.24.0 → 0.26.0

package/src/memory.ts

@@ -8,6 +8,7 @@ import {
   slugifyStorageComponent,
   STORAGE_SCHEMA_VERSION,
 } from "./agent-identity.js";
+import { createRawKVStore, type RawKVStore } from "./kv-store.js";
 
 export interface MainMemory {
   content: string;
@@ -27,7 +28,7 @@ export interface MemoryConfig {
 
 export interface MemoryStore {
   getMainMemory(): Promise<MainMemory>;
-  updateMainMemory(input: { content: string; mode?: "replace" | "append" }): Promise<MainMemory>;
+  updateMainMemory(input: { content: string }): Promise<MainMemory>;
 }
 
 type MainMemoryPayload = {
@@ -89,19 +90,10 @@ class InMemoryMemoryStore implements MemoryStore {
     return this.mainMemory;
   }
 
-  async updateMainMemory(input: {
-    content: string;
-    mode?: "replace" | "append";
-  }): Promise<MainMemory> {
-    const now = Date.now();
-    const existing = await this.getMainMemory();
-    const nextContent =
-      input.mode === "append" && existing.content
-        ? `${existing.content}\n\n${input.content}`.trim()
-        : input.content;
+  async updateMainMemory(input: { content: string }): Promise<MainMemory> {
     this.mainMemory = {
-      content: nextContent.trim(),
-      updatedAt: now,
+      content: input.content.trim(),
+      updatedAt: Date.now(),
     };
     return this.mainMemory;
   }
@@ -166,18 +158,10 @@ class FileMainMemoryStore implements MemoryStore {
     return this.mainMemory;
   }
 
-  async updateMainMemory(input: {
-    content: string;
-    mode?: "replace" | "append";
-  }): Promise<MainMemory> {
+  async updateMainMemory(input: { content: string }): Promise<MainMemory> {
     await this.ensureLoaded();
-    const existing = await this.getMainMemory();
-    const nextContent =
-      input.mode === "append" && existing.content
-        ? `${existing.content}\n\n${input.content}`.trim()
-        : input.content;
     this.mainMemory = {
-      content: nextContent.trim(),
+      content: input.content.trim(),
       updatedAt: Date.now(),
     };
     await this.persist();
@@ -185,25 +169,23 @@ class FileMainMemoryStore implements MemoryStore {
   }
 }
 
-abstract class KeyValueMainMemoryStoreBase implements MemoryStore {
-  protected readonly ttl?: number;
-  protected readonly memoryFallback: InMemoryMemoryStore;
+class KVBackedMemoryStore implements MemoryStore {
+  private readonly kv: RawKVStore;
+  private readonly storageKey: string;
+  private readonly ttl?: number;
+  private readonly memoryFallback: InMemoryMemoryStore;
 
-  constructor(ttl?: number) {
+  constructor(kv: RawKVStore, storageKey: string, ttl?: number) {
+    this.kv = kv;
+    this.storageKey = storageKey;
     this.ttl = ttl;
     this.memoryFallback = new InMemoryMemoryStore(ttl);
   }
 
-  protected abstract getRaw(key: string): Promise<string | undefined>;
-  protected abstract setRaw(key: string, value: string): Promise<void>;
-  protected abstract setRawWithTtl(key: string, value: string, ttl: number): Promise<void>;
-
-  protected async readPayload(key: string): Promise<MainMemoryPayload> {
+  private async readPayload(): Promise<MainMemoryPayload> {
     try {
-      const raw = await this.getRaw(key);
-      if (!raw) {
-        return { main: { ...DEFAULT_MAIN_MEMORY } };
-      }
+      const raw = await this.kv.get(this.storageKey);
+      if (!raw) return { main: { ...DEFAULT_MAIN_MEMORY } };
       const parsed = JSON.parse(raw) as MainMemoryPayload;
       const content = typeof parsed.main?.content === "string" ? parsed.main.content : "";
       const updatedAt = typeof parsed.main?.updatedAt === "number" ? parsed.main.updatedAt : 0;
@@ -214,267 +196,30 @@ abstract class KeyValueMainMemoryStoreBase implements MemoryStore {
     }
   }
 
-  protected async writePayload(key: string, payload: MainMemoryPayload): Promise<void> {
+  private async writePayload(payload: MainMemoryPayload): Promise<void> {
     try {
       const serialized = JSON.stringify(payload);
       if (typeof this.ttl === "number") {
-        await this.setRawWithTtl(key, serialized, Math.max(1, this.ttl));
+        await this.kv.setWithTtl(this.storageKey, serialized, Math.max(1, this.ttl));
       } else {
-        await this.setRaw(key, serialized);
+        await this.kv.set(this.storageKey, serialized);
       }
     } catch {
-      await this.memoryFallback.updateMainMemory({
-        content: payload.main.content,
-        mode: "replace",
-      });
+      await this.memoryFallback.updateMainMemory({ content: payload.main.content });
     }
   }
 
   async getMainMemory(): Promise<MainMemory> {
-    const payload = await this.readPayload(this.key());
+    const payload = await this.readPayload();
     return payload.main;
   }
 
-  async updateMainMemory(input: {
-    content: string;
-    mode?: "replace" | "append";
-  }): Promise<MainMemory> {
-    const key = this.key();
-    const payload = await this.readPayload(key);
-    const nextContent =
-      input.mode === "append" && payload.main.content
-        ? `${payload.main.content}\n\n${input.content}`.trim()
-        : input.content;
-    payload.main = {
-      content: nextContent.trim(),
-      updatedAt: Date.now(),
-    };
-    await this.writePayload(key, payload);
+  async updateMainMemory(input: { content: string }): Promise<MainMemory> {
+    const payload = await this.readPayload();
+    payload.main = { content: input.content.trim(), updatedAt: Date.now() };
+    await this.writePayload(payload);
     return payload.main;
   }
-
-  protected abstract key(): string;
-}
-
-class UpstashMemoryStore extends KeyValueMainMemoryStoreBase {
-  private readonly baseUrl: string;
-  private readonly token: string;
-  private readonly storageKey: string;
-
-  constructor(options: {
-    baseUrl: string;
-    token: string;
-    storageKey: string;
-    ttl?: number;
-  }) {
-    super(options.ttl);
-    this.baseUrl = options.baseUrl.replace(/\/+$/, "");
-    this.token = options.token;
-    this.storageKey = options.storageKey;
-  }
-
-  protected key(): string {
-    return this.storageKey;
-  }
-
-  private headers(): HeadersInit {
-    return {
-      Authorization: `Bearer ${this.token}`,
-      "Content-Type": "application/json",
-    };
-  }
-
-  protected async getRaw(key: string): Promise<string | undefined> {
-    const response = await fetch(`${this.baseUrl}/get/${encodeURIComponent(key)}`, {
-      method: "POST",
-      headers: this.headers(),
-    });
-    if (!response.ok) {
-      return undefined;
-    }
-    const payload = (await response.json()) as { result?: string | null };
-    return payload.result ?? undefined;
-  }
-
-  protected async setRaw(key: string, value: string): Promise<void> {
-    await fetch(
-      `${this.baseUrl}/set/${encodeURIComponent(key)}/${encodeURIComponent(value)}`,
-      { method: "POST", headers: this.headers() },
-    );
-  }
-
-  protected async setRawWithTtl(key: string, value: string, ttl: number): Promise<void> {
-    await fetch(
-      `${this.baseUrl}/setex/${encodeURIComponent(key)}/${Math.max(1, ttl)}/${encodeURIComponent(
-        value,
-      )}`,
-      { method: "POST", headers: this.headers() },
-    );
-  }
-}
-
-class RedisMemoryStore extends KeyValueMainMemoryStoreBase {
-  private readonly storageKey: string;
-  private readonly clientPromise: Promise<
-    | {
-        get: (key: string) => Promise<string | null>;
-        set: (key: string, value: string, options?: { EX?: number }) => Promise<unknown>;
-      }
-    | undefined
-  >;
-
-  constructor(options: {
-    url: string;
-    storageKey: string;
-    ttl?: number;
-  }) {
-    super(options.ttl);
-    this.storageKey = options.storageKey;
-    this.clientPromise = (async () => {
-      try {
-        const redisModule = (await import("redis")) as unknown as {
-          createClient: (args: { url: string }) => {
-            connect: () => Promise<unknown>;
-            get: (key: string) => Promise<string | null>;
-            set: (key: string, value: string, options?: { EX?: number }) => Promise<unknown>;
-          };
-        };
-        const client = redisModule.createClient({ url: options.url });
-        await client.connect();
-        return client;
-      } catch {
-        return undefined;
-      }
-    })();
-  }
-
-  protected key(): string {
-    return this.storageKey;
-  }
-
-  protected async getRaw(key: string): Promise<string | undefined> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("Redis unavailable");
-    }
-    const value = await client.get(key);
-    return value ?? undefined;
-  }
-
-  protected async setRaw(key: string, value: string): Promise<void> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("Redis unavailable");
-    }
-    await client.set(key, value);
-  }
-
-  protected async setRawWithTtl(key: string, value: string, ttl: number): Promise<void> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("Redis unavailable");
-    }
-    await client.set(key, value, { EX: Math.max(1, ttl) });
-  }
-}
-
-class DynamoDbMemoryStore extends KeyValueMainMemoryStoreBase {
-  private readonly storageKey: string;
-  private readonly table: string;
-  private readonly clientPromise: Promise<
-    | {
-        send: (command: unknown) => Promise<unknown>;
-        GetItemCommand: new (input: unknown) => unknown;
-        PutItemCommand: new (input: unknown) => unknown;
-      }
-    | undefined
-  >;
-
-  constructor(options: {
-    table: string;
-    storageKey: string;
-    region?: string;
-    ttl?: number;
-  }) {
-    super(options.ttl);
-    this.storageKey = options.storageKey;
-    this.table = options.table;
-    this.clientPromise = (async () => {
-      try {
-        const module = (await import("@aws-sdk/client-dynamodb")) as {
-          DynamoDBClient: new (input: { region?: string }) => {
-            send: (command: unknown) => Promise<unknown>;
-          };
-          GetItemCommand: new (input: unknown) => unknown;
-          PutItemCommand: new (input: unknown) => unknown;
-        };
-        const client = new module.DynamoDBClient({ region: options.region });
-        return {
-          send: client.send.bind(client),
-          GetItemCommand: module.GetItemCommand,
-          PutItemCommand: module.PutItemCommand,
-        };
-      } catch {
-        return undefined;
-      }
-    })();
-  }
-
-  protected key(): string {
-    return this.storageKey;
-  }
-
-  protected async getRaw(key: string): Promise<string | undefined> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("DynamoDB unavailable");
-    }
-    const result = (await client.send(
-      new client.GetItemCommand({
-        TableName: this.table,
-        Key: { runId: { S: key } },
-      }),
-    )) as {
-      Item?: {
-        value?: { S?: string };
-      };
-    };
-    return result.Item?.value?.S;
-  }
-
-  protected async setRaw(key: string, value: string): Promise<void> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("DynamoDB unavailable");
-    }
-    await client.send(
-      new client.PutItemCommand({
-        TableName: this.table,
-        Item: {
-          runId: { S: key },
-          value: { S: value },
-        },
-      }),
-    );
-  }
-
-  protected async setRawWithTtl(key: string, value: string, ttl: number): Promise<void> {
-    const client = await this.clientPromise;
-    if (!client) {
-      throw new Error("DynamoDB unavailable");
-    }
-    const ttlEpoch = Math.floor(Date.now() / 1000) + Math.max(1, ttl);
-    await client.send(
-      new client.PutItemCommand({
-        TableName: this.table,
-        Item: {
-          runId: { S: key },
-          value: { S: value },
-          ttl: { N: String(ttlEpoch) },
-        },
-      }),
-    );
-  }
 }
 
 export const createMemoryStore = (
@@ -484,54 +229,19 @@ export const createMemoryStore = (
 ): MemoryStore => {
   const provider = config?.provider ?? "local";
   const ttl = config?.ttl;
-  const storageKey = `poncho:${STORAGE_SCHEMA_VERSION}:${slugifyStorageComponent(
-    agentId,
-  )}:memory:main`;
   const workingDir = options?.workingDir ?? process.cwd();
+
   if (provider === "local") {
     return new FileMainMemoryStore(workingDir, ttl);
   }
   if (provider === "memory") {
     return new InMemoryMemoryStore(ttl);
   }
-  if (provider === "upstash") {
-    const urlEnv = config?.urlEnv ?? (process.env.UPSTASH_REDIS_REST_URL ? "UPSTASH_REDIS_REST_URL" : "KV_REST_API_URL");
-    const tokenEnv = config?.tokenEnv ?? (process.env.UPSTASH_REDIS_REST_TOKEN ? "UPSTASH_REDIS_REST_TOKEN" : "KV_REST_API_TOKEN");
-    const url = process.env[urlEnv] ?? "";
-    const token = process.env[tokenEnv] ?? "";
-    if (url && token) {
-      return new UpstashMemoryStore({
-        baseUrl: url,
-        token,
-        storageKey,
-        ttl,
-      });
-    }
-    return new InMemoryMemoryStore(ttl);
-  }
-  if (provider === "redis") {
-    const urlEnv = config?.urlEnv ?? "REDIS_URL";
-    const url = process.env[urlEnv] ?? "";
-    if (url) {
-      return new RedisMemoryStore({
-        url,
-        storageKey,
-        ttl,
-      });
-    }
-    return new InMemoryMemoryStore(ttl);
-  }
-  if (provider === "dynamodb") {
-    const table = config?.table ?? process.env.PONCHO_DYNAMODB_TABLE ?? "";
-    if (table) {
-      return new DynamoDbMemoryStore({
-        table,
-        storageKey,
-        region: config?.region,
-        ttl,
-      });
-    }
-    return new InMemoryMemoryStore(ttl);
+
+  const kv = createRawKVStore(config);
+  if (kv) {
+    const storageKey = `poncho:${STORAGE_SCHEMA_VERSION}:${slugifyStorageComponent(agentId)}:memory:main`;
+    return new KVBackedMemoryStore(kv, storageKey, ttl);
   }
   return new InMemoryMemoryStore(ttl);
 };
@@ -590,20 +300,17 @@ export const createMemoryTools = (
       },
     }),
     defineTool({
-      name: "memory_main_update",
+      name: "memory_main_write",
       description:
-        "Update persistent main memory when new stable preferences, long-term goals, or durable facts appear. Proactively evaluate every turn whether memory should be updated, and avoid storing ephemeral details.",
+        "Overwrite the entire persistent main memory document. " +
+        "Use for initial writes or full rewrites. " +
+        "Prefer memory_main_edit for targeted changes to existing memory.",
       inputSchema: {
         type: "object",
         properties: {
-          mode: {
-            type: "string",
-            enum: ["replace", "append"],
-            description: "replace overwrites memory; append adds content to the end",
-          },
           content: {
             type: "string",
-            description: "The memory content to write",
+            description: "The full memory content to write",
           },
         },
         required: ["content"],
@@ -614,11 +321,56 @@ export const createMemoryTools = (
         if (!content) {
           throw new Error("content is required");
         }
-        const mode =
-          input.mode === "append" || input.mode === "replace"
-            ? input.mode
-            : "replace";
-        const memory = await store.updateMainMemory({ content, mode });
+        const memory = await store.updateMainMemory({ content });
+        return { ok: true, memory };
+      },
+    }),
+    defineTool({
+      name: "memory_main_edit",
+      description:
+        "Edit persistent main memory by replacing an exact string match with new content. " +
+        "The old_str must match exactly one location in memory. " +
+        "Use an empty new_str to delete matched content. " +
+        "Proactively evaluate every turn whether memory should be updated.",
+      inputSchema: {
+        type: "object",
+        properties: {
+          old_str: {
+            type: "string",
+            description:
+              "The exact text to find and replace (must be unique in memory). " +
+              "Include surrounding context if needed to ensure uniqueness.",
+          },
+          new_str: {
+            type: "string",
+            description: "The replacement text (use empty string to delete the matched content)",
+          },
+        },
+        required: ["old_str", "new_str"],
+        additionalProperties: false,
+      },
+      handler: async (input) => {
+        const oldStr = typeof input.old_str === "string" ? input.old_str : "";
+        const newStr = typeof input.new_str === "string" ? input.new_str : "";
+        if (!oldStr) {
+          throw new Error("old_str must not be empty.");
+        }
+        const current = await store.getMainMemory();
+        const content = current.content;
+        const first = content.indexOf(oldStr);
+        if (first === -1) {
+          throw new Error(
+            "old_str not found in memory. Make sure it matches exactly, including whitespace and line breaks.",
+          );
+        }
+        const last = content.lastIndexOf(oldStr);
+        if (first !== last) {
+          throw new Error(
+            "old_str appears multiple times in memory. Please provide more context to ensure a unique match.",
+          );
+        }
+        const newContent = content.slice(0, first) + newStr + content.slice(first + oldStr.length);
+        const memory = await store.updateMainMemory({ content: newContent });
         return { ok: true, memory };
       },
     }),