feat: persist pruned tool IDs across OpenCode restarts

Adds state persistence so pruned tool IDs survive OpenCode restarts.
State is saved per-session to ~/.local/share/opencode/storage/plugin/dcp/{sessionId}.json

Closes #16

Author: Tarquinen

lib/fetch-wrapper/gemini.ts

@@ -53,7 +53,7 @@ export async function handleGemini(
         return { modified, body }
     }
 
-    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state)
+    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger)
 
     if (allPrunedIds.size === 0) {
         return { modified, body }

lib/fetch-wrapper/openai-chat.ts

@@ -61,7 +61,7 @@ export async function handleOpenAIChatAndAnthropic(
         return false
     })
 
-    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state)
+    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger)
 
     if (toolMessages.length === 0 || allPrunedIds.size === 0) {
         return { modified, body }

lib/fetch-wrapper/openai-responses.ts

@@ -55,7 +55,7 @@ export async function handleOpenAIResponses(
         return { modified, body }
     }
 
-    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state)
+    const { allSessions, allPrunedIds } = await getAllPrunedIds(ctx.client, ctx.state, ctx.logger)
 
     if (allPrunedIds.size === 0) {
         return { modified, body }

lib/fetch-wrapper/types.ts

@@ -1,4 +1,4 @@
-import type { PluginState } from "../state"
+import { type PluginState, ensureSessionRestored } from "../state"
 import type { Logger } from "../logger"
 import type { ToolTracker } from "../synth-instruction"
 import type { PluginConfig } from "../config"
@@ -36,22 +36,19 @@ export interface PrunedIdData {
     allPrunedIds: Set<string>
 }
 
-/**
- * Get all pruned IDs across all non-subagent sessions.
- */
 export async function getAllPrunedIds(
     client: any,
-    state: PluginState
+    state: PluginState,
+    logger?: Logger
 ): Promise<PrunedIdData> {
     const allSessions = await client.session.list()
     const allPrunedIds = new Set<string>()
 
-    if (allSessions.data) {
-        for (const session of allSessions.data) {
-            if (session.parentID) continue
-            const prunedIds = state.prunedIds.get(session.id) ?? []
-            prunedIds.forEach((id: string) => allPrunedIds.add(id))
-        }
+    const currentSession = getMostRecentActiveSession(allSessions)
+    if (currentSession) {
+        await ensureSessionRestored(state, currentSession.id, logger)
+        const prunedIds = state.prunedIds.get(currentSession.id) ?? []
+        prunedIds.forEach((id: string) => allPrunedIds.add(id))
     }
 
     return { allSessions, allPrunedIds }

lib/janitor.ts

@@ -6,6 +6,7 @@ import { selectModel, extractModelFromSession } from "./model-selector"
 import { estimateTokensBatch, formatTokenCount } from "./tokenizer"
 import { detectDuplicates } from "./deduplicator"
 import { extractParameterKey } from "./display-utils"
+import { saveSessionState } from "./state-persistence"
 
 export interface SessionStats {
     totalToolsPruned: number
@@ -97,8 +98,6 @@ export class Janitor {
                 return null
             }
 
-            // Extract the current agent from the last user message to preserve agent context
-            // Following the same pattern as OpenCode's server.ts
             let currentAgent: string | undefined = undefined
             for (let i = messages.length - 1; i >= 0; i--) {
                 const msg = messages[i]
@@ -330,6 +329,10 @@ export class Janitor {
             const allPrunedIds = [...new Set([...alreadyPrunedIds, ...finalPrunedIds])]
             this.prunedIdsState.set(sessionID, allPrunedIds)
 
+            saveSessionState(sessionID, new Set(allPrunedIds), sessionStats, this.logger).catch(err => {
+                this.logger.error("janitor", "Failed to persist state", { error: err.message })
+            })
+
             const prunedCount = finalNewlyPrunedIds.length
             const keptCount = candidateCount - prunedCount
             const hasBoth = deduplicatedIds.length > 0 && llmPrunedIds.length > 0

lib/state-persistence.ts

@@ -0,0 +1,111 @@
+/**
+ * State persistence module for DCP plugin.
+ * Persists pruned tool IDs across sessions so they survive OpenCode restarts.
+ * Storage location: ~/.local/share/opencode/storage/plugin/dcp/{sessionId}.json
+ */
+
+import * as fs from "fs/promises";
+import { existsSync } from "fs";
+import { homedir } from "os";
+import { join } from "path";
+import type { SessionStats } from "./janitor";
+import type { Logger } from "./logger";
+
+export interface PersistedSessionState {
+    prunedIds: string[];
+    stats: SessionStats;
+    lastUpdated: string;
+}
+
+const STORAGE_DIR = join(
+    homedir(),
+    ".local",
+    "share",
+    "opencode",
+    "storage",
+    "plugin",
+    "dcp"
+);
+
+async function ensureStorageDir(): Promise<void> {
+    if (!existsSync(STORAGE_DIR)) {
+        await fs.mkdir(STORAGE_DIR, { recursive: true });
+    }
+}
+
+function getSessionFilePath(sessionId: string): string {
+    return join(STORAGE_DIR, `${sessionId}.json`);
+}
+
+export async function saveSessionState(
+    sessionId: string,
+    prunedIds: Set<string>,
+    stats: SessionStats,
+    logger?: Logger
+): Promise<void> {
+    try {
+        await ensureStorageDir();
+
+        const state: PersistedSessionState = {
+            prunedIds: Array.from(prunedIds),
+            stats,
+            lastUpdated: new Date().toISOString(),
+        };
+
+        const filePath = getSessionFilePath(sessionId);
+        const content = JSON.stringify(state, null, 2);
+        await fs.writeFile(filePath, content, "utf-8");
+
+        logger?.info("persist", "Saved session state to disk", {
+            sessionId: sessionId.slice(0, 8),
+            prunedIds: prunedIds.size,
+            totalTokensSaved: stats.totalTokensSaved,
+        });
+    } catch (error: any) {
+        logger?.error("persist", "Failed to save session state", {
+            sessionId: sessionId.slice(0, 8),
+            error: error?.message,
+        });
+    }
+}
+
+export async function loadSessionState(
+    sessionId: string,
+    logger?: Logger
+): Promise<PersistedSessionState | null> {
+    try {
+        const filePath = getSessionFilePath(sessionId);
+
+        if (!existsSync(filePath)) {
+            return null;
+        }
+
+        const content = await fs.readFile(filePath, "utf-8");
+        const state = JSON.parse(content) as PersistedSessionState;
+
+        if (!state || !Array.isArray(state.prunedIds) || !state.stats) {
+            logger?.warn("persist", "Invalid session state file, ignoring", {
+                sessionId: sessionId.slice(0, 8),
+            });
+            return null;
+        }
+
+        logger?.info("persist", "Loaded session state from disk", {
+            sessionId: sessionId.slice(0, 8),
+            prunedIds: state.prunedIds.length,
+            totalTokensSaved: state.stats.totalTokensSaved,
+        });
+
+        return state;
+    } catch (error: any) {
+        logger?.warn("persist", "Failed to load session state", {
+            sessionId: sessionId.slice(0, 8),
+            error: error?.message,
+        });
+        return null;
+    }
+}
+
+export function getStorageDir(): string {
+    return STORAGE_DIR;
+}

lib/state.ts

@@ -1,4 +1,6 @@
 import type { SessionStats } from "./janitor"
+import type { Logger } from "./logger"
+import { loadSessionState } from "./state-persistence"
 
 /**
  * Centralized state management for the DCP plugin.
@@ -18,6 +20,8 @@ export interface PluginState {
      * Key: sessionID, Value: Map<positionKey, toolCallId> where positionKey is "toolName:index"
      */
     googleToolCallMapping: Map<string, Map<string, string>>
+    /** Set of session IDs that have been restored from disk */
+    restoredSessions: Set<string>
 }
 
 export interface ToolParameterEntry {
@@ -40,5 +44,32 @@ export function createPluginState(): PluginState {
         toolParameters: new Map(),
         model: new Map(),
         googleToolCallMapping: new Map(),
+        restoredSessions: new Set(),
+    }
+}
+
+export async function ensureSessionRestored(
+    state: PluginState,
+    sessionId: string,
+    logger?: Logger
+): Promise<void> {
+    if (state.restoredSessions.has(sessionId)) {
+        return
+    }
+
+    state.restoredSessions.add(sessionId)
+
+    const persisted = await loadSessionState(sessionId, logger)
+    if (persisted) {
+        if (!state.prunedIds.has(sessionId)) {
+            state.prunedIds.set(sessionId, persisted.prunedIds)
+            logger?.info("persist", "Restored prunedIds from disk", {
+                sessionId: sessionId.slice(0, 8),
+                count: persisted.prunedIds.length,
+            })
+        }
+        if (!state.stats.has(sessionId)) {
+            state.stats.set(sessionId, persisted.stats)
+        }
     }
 }