The AI Agent Split-Brain Trap: Why Static Indexes Destroy Long-Running Context (And How to Build Transactional Memory)

// Transactional Agent Memory Engine for Cloud Run and Cloud Firestore
import { initializeApp, getApps } from 'firebase-admin/app';
import { getFirestore, 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 interface MutationResult {
  success: boolean;
  newVersion?: number;
  currentState?: AgentTaskResponse;
  error?: string;
}

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 || {},
    };
  }

  /**
   * Updates an agent task using optimistic concurrency control (OCC).
   * Prevents reasoning staleness across multi-turn LLM planning cycles.
   */
  async updateTaskAtomic(
    taskId: string,
    expectedVersion: number,
    updates: Partial<Pick<AgentTaskDocument, 'status' | 'assignedAgent' | 'payload'>>
  ): Promise<MutationResult> {
    const taskRef = this.tasksCol.doc(taskId);

    try {
      const result = await db.runTransaction(async (transaction) => {
        const snapshot = await transaction.get(taskRef);

        if (!snapshot.exists) {
          throw new Error(`Task ${taskId} does not exist.`);
        }

        const currentData = snapshot.data() as AgentTaskDocument;

        if (currentData.version !== expectedVersion) {
          const conflictError = new Error('VERSION_CONFLICT');
          (conflictError as any).currentState = this.formatTask(taskId, currentData);
          throw conflictError;
        }

        const newVersion = currentData.version + 1;

        transaction.update(taskRef, {
          ...updates,
          version: newVersion,
          lastUpdated: FieldValue.serverTimestamp(),
        });

        return newVersion;
      });

      return { success: true, newVersion: result };
    } catch (err: unknown) {
      if (err instanceof Error && err.message === 'VERSION_CONFLICT') {
        const conflict = err as any;
        return {
          success: false,
          error: `Concurrency collision on task ${taskId}. Version advanced before write.`,
          currentState: conflict.currentState,
        };
      }

      const message = err instanceof Error ? err.message : 'Unknown transaction failure';
      return { success: false, error: message };
    }
  }

  /**
   * 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" }
      ]
    }
  ]
}