Two Writers, One Index: How Static Files Corrupt Agent Memory

// Transactional Agent Memory Engine for Cloud Run and Cloud Firestore
import { initializeApp, getApps } from 'firebase-admin/app';
import { getFirestore, FieldPath, FieldValue, Timestamp } from 'firebase-admin/firestore';

if (getApps().length === 0) {
  initializeApp(); // Uses Application Default Credentials on Cloud Run
}

const db = getFirestore();

export type TaskStatus = 'PENDING' | 'IN_PROGRESS' | 'COMPLETED' | 'FAILED';

export interface AgentTaskDocument {
  title: string;
  status: TaskStatus;
  version: number;
  assignedAgent: string;
  lastUpdated: Timestamp;
  payload: Record<string, unknown>;
}

export interface AgentTaskResponse {
  id: string;
  title: string;
  status: TaskStatus;
  version: number;
  assignedAgent: string;
  lastUpdated: string;
  payload: Record<string, unknown>;
}

export type TaskUpdates = Partial<Pick<AgentTaskDocument, 'status' | 'assignedAgent' | 'payload'>>;

export type MutationResult =
  | { success: true; newVersion: number }
  | { success: false; reason: 'NOT_FOUND' }
  | { success: false; reason: 'VERSION_CONFLICT'; currentState: AgentTaskResponse };

export class TransactionalMemoryEngine {
  private tasksCol = db.collection('agent_tasks');

  private formatTask(id: string, data: AgentTaskDocument): AgentTaskResponse {
    return {
      id,
      title: data.title,
      status: data.status,
      version: data.version,
      assignedAgent: data.assignedAgent,
      lastUpdated: data.lastUpdated
        ? data.lastUpdated.toDate().toISOString()
        : new Date().toISOString(),
      payload: data.payload || {},
    };
  }

  /**
   * Writes only if nobody changed the task since the caller read expectedVersion.
   * The version field is the optimistic check; the transaction makes read, check, and write atomic.
   */
  async updateTaskAtomic(
    taskId: string,
    expectedVersion: number,
    updates: TaskUpdates
  ): Promise<MutationResult> {
    const taskRef = this.tasksCol.doc(taskId);
    const { payload = {}, ...fields } = updates;

    // Expected outcomes are return values; anything thrown (permissions, deadlines, contention) propagates
    return db.runTransaction(async (transaction): Promise<MutationResult> => {
      const snapshot = await transaction.get(taskRef);
      if (!snapshot.exists) {
        return { success: false, reason: 'NOT_FOUND' };
      }

      const currentData = snapshot.data() as AgentTaskDocument;
      if (currentData.version !== expectedVersion) {
        return { success: false, reason: 'VERSION_CONFLICT', currentState: this.formatTask(taskId, currentData) };
      }

      const newVersion = currentData.version + 1;
      // Field/value pairs: a FieldPath merges each key into the payload map as-is, even with dots or slashes
      transaction.update(taskRef, 'version', newVersion, 'lastUpdated', FieldValue.serverTimestamp(),
        ...Object.entries(fields).flat(),
        ...Object.entries(payload).flatMap(([key, value]) => [new FieldPath('payload', key), value]));
      return { success: true, newVersion };
    });
  }

  /**
   * On a version conflict, re-plans against the fresh state and retries, as agent B does in Figure 1.
   */
  async updateWithRetry(
    task: AgentTaskResponse,
    plan: (current: AgentTaskResponse) => TaskUpdates,
    maxAttempts = 3
  ): Promise<MutationResult> {
    let current = task;
    for (let attempt = 1; ; attempt++) {
      const result = await this.updateTaskAtomic(current.id, current.version, plan(current));
      if (result.success || result.reason !== 'VERSION_CONFLICT' || attempt >= maxAttempts) return result;
      current = result.currentState; // re-read: plan again on top of the write that won
    }
  }

  /**
   * Fetches fresh, query-indexed state directly from Firestore.
   */
  async getActiveTasksForAgent(
    agentId: string,
    limitCount = 10
  ): Promise<AgentTaskResponse[]> {
    const querySnapshot = await this.tasksCol
      .where('assignedAgent', '==', agentId)
      .where('status', 'in', ['PENDING', 'IN_PROGRESS'])
      .orderBy('lastUpdated', 'desc')
      .limit(limitCount)
      .get();

    return querySnapshot.docs.map((doc) =>
      this.formatTask(doc.id, doc.data() as AgentTaskDocument)
    );
  }
}

{
  "indexes": [
    {
      "collectionGroup": "agent_tasks",
      "queryScope": "COLLECTION",
      "fields": [
        { "fieldPath": "assignedAgent", "order": "ASCENDING" },
        { "fieldPath": "status", "order": "ASCENDING" },
        { "fieldPath": "lastUpdated", "order": "DESCENDING" }
      ]
    }
  ]
}