src/agent/SupabaseSaver.ts

import { 
  BaseCheckpointSaver, 
  type Checkpoint, 
  type CheckpointListOptions, 
  type CheckpointTuple, 
  type PendingWrite, 
  type CheckpointMetadata,
  type ChannelVersions
} from '@langchain/langgraph-checkpoint';
import type { RunnableConfig } from '@langchain/core/runnables';
import { getSupabaseClient } from '../supabase.js';

export class SupabaseSaver extends BaseCheckpointSaver {
  async getTuple(config: RunnableConfig): Promise<CheckpointTuple | undefined> {
    const threadId = config.configurable?.thread_id;
    if (!threadId) return undefined;
    
    const client = getSupabaseClient();
    const query = client.from('langgraph_checkpoints')
      .select('*')
      .eq('thread_id', threadId)
      .order('checkpoint_id', { ascending: false })
      .limit(1);
    
    if (config.configurable?.checkpoint_id) {
      query.eq('checkpoint_id', config.configurable.checkpoint_id);
    }
    
    const { data, error } = await query.single();
    if (error || !data) return undefined;
    
    const checkpoint = await this.serde.loadsTyped('json', data.checkpoint);
    const metadata = await this.serde.loadsTyped('json', data.metadata);
    
    return {
      config: {
        configurable: {
          thread_id: threadId,
          checkpoint_ns: '',
          checkpoint_id: data.checkpoint_id
        }
      },
      checkpoint: checkpoint as Checkpoint,
      metadata: metadata as CheckpointMetadata,
      parentConfig: data.parent_checkpoint_id ? {
        configurable: {
          thread_id: threadId,
          checkpoint_ns: '',
          checkpoint_id: data.parent_checkpoint_id
        }
      } : undefined
    };
  }

  async *list(config: RunnableConfig, options?: CheckpointListOptions): AsyncGenerator<CheckpointTuple> {
    const threadId = config.configurable?.thread_id;
    if (!threadId) return;
    
    const client = getSupabaseClient();
    let query = client.from('langgraph_checkpoints')
      .select('*')
      .eq('thread_id', threadId)
      .order('checkpoint_id', { ascending: false });
    
    if (options?.limit) {
      query = query.limit(options.limit);
    }
    
    const { data, error } = await query;
    if (error || !data) return;
    
    for (const row of data) {
      const checkpoint = await this.serde.loadsTyped('json', row.checkpoint);
      const metadata = await this.serde.loadsTyped('json', row.metadata);
      yield {
        config: { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: row.checkpoint_id } },
        checkpoint: checkpoint as Checkpoint,
        metadata: metadata as CheckpointMetadata,
        parentConfig: row.parent_checkpoint_id ? { configurable: { thread_id: threadId, checkpoint_ns: '', checkpoint_id: row.parent_checkpoint_id } } : undefined
      };
    }
  }

  async put(
    config: RunnableConfig, 
    checkpoint: Checkpoint, 
    metadata: CheckpointMetadata,
    newVersions: ChannelVersions
  ): Promise<RunnableConfig> {
    const threadId = config.configurable?.thread_id;
    if (!threadId) throw new Error('missing thread_id');
    
    const [ckptStr, metaStr] = await Promise.all([
      this.serde.dumpsTyped(checkpoint),
      this.serde.dumpsTyped(metadata)
    ]);
    
    const client = getSupabaseClient();
    const { error } = await client.from('langgraph_checkpoints').upsert({
      thread_id: threadId,
      checkpoint_id: checkpoint.id,
      parent_checkpoint_id: config.configurable?.checkpoint_id || null,
      checkpoint: ckptStr[1],
      metadata: metaStr[1]
    });
    
    if (error) {
      throw new Error(`Failed to put checkpoint: ${error.message}`);
    }
    
    return {
      configurable: {
        thread_id: threadId,
        checkpoint_ns: '',
        checkpoint_id: checkpoint.id
      }
    };
  }

  async putWrites(config: RunnableConfig, writes: PendingWrite[], taskId: string): Promise<void> {
    const threadId = config.configurable?.thread_id;
    const checkpointId = config.configurable?.checkpoint_id;
    if (!threadId || !checkpointId) throw new Error('missing thread_id or checkpoint_id');
    
    const client = getSupabaseClient();
    const rows = await Promise.all(writes.map(async (w, idx) => {
      const [type, value] = await this.serde.dumpsTyped(w[1]);
      return {
        thread_id: threadId,
        checkpoint_id: checkpointId,
        task_id: taskId,
        idx,
        channel: w[0],
        type,
        value
      };
    }));
    
    if (rows.length > 0) {
      const { error } = await client.from('langgraph_checkpoint_writes').upsert(rows);
      if (error) {
        throw new Error(`Failed to put writes: ${error.message}`);
      }
    }
  }

  async deleteThread(threadId: string): Promise<void> {
    const client = getSupabaseClient();
    await client.from('langgraph_checkpoint_writes').delete().eq('thread_id', threadId);
    await client.from('langgraph_checkpoints').delete().eq('thread_id', threadId);
  }
}