@amaster.ai/dsh-storage 0.1.33 → 0.1.34

package/README.md

@@ -55,7 +55,7 @@ The logical message id rides in `metadata.id`; message rows use a deterministic
 
 For `user/message` rows, `content` remains the text projection, while `metadata.message` retains the original user message as received, including `source` and ordered content blocks. Image/file blocks keep whatever attachment reference, path, URL, or inline data the incoming message actually contains; this mirror does not fetch or rewrite attachments. Re-projection writes the same raw message into the existing row.
 
-The `tool_calls` array retains dsh's tool-call blocks (`id`, `name`, `arguments`); the matching entry gains the structured `result` block and, on failure, an `error` identity. The plugin reads persisted model rows to match results across process restarts and upserts the same row ID, so duplicate results do not add rows or increment the session's message count. Existing historical `type=tool` rows are left intact.
+The `tool_calls` array retains dsh's tool-call blocks (`id`, `name`, `arguments`); the matching entry gains the structured `result` block and, on failure, an `error` identity. Live call IDs are cached, and restored calls can resolve their message ID from the model surface: both paths read one message primary key. Older calls are searched in keyset pages of at most 100 ID/tool-call candidates, matched in the plugin and discarded before fetching the next page. Tool-result lookup SQL does not search JSON, join tables, or nest SELECT statements. Duplicate results update the same row and do not increment message counts. Existing historical `type=tool` rows are left intact.
 
 When an assistant message contains dsh `reasoning` blocks, `thoughts` is an array of `{ content }` entries in block order. If the adapter supplies a matching signature in the message's `source.replayState.blocks`, that entry also carries `thoughtSignature`. The full model `source`, including opaque `replayState`, is preserved in row metadata for lossless provider replay. The backend's `readMessages()` also accepts older `{ description }` entries and JSON string thoughts, returning them as `{ content }`. `subject` and a per-thought timestamp are not invented when the upstream message does not supply them.
 

package/lib/backends/database.js

@@ -143,18 +143,89 @@ export class DatabaseBackend {
         });
     }
     async readMessages(sessionId, messageId) {
+        return this.readRows({
+            sessionId,
+            deletedAt: null,
+            id: pk(`message ${sessionId}\0${messageId}`),
+        });
+    }
+    async readToolCallMessages(sessionId, callId) {
+        if (!this.prisma)
+            return [];
+        const provider = this.config.provider;
+        const sqlserver = provider === 'sqlserver';
+        const sortTime = provider === 'mysql'
+            ? "DATE_FORMAT(created_at, '%Y-%m-%d %H:%i:%s.%f')"
+            : provider === 'postgresql'
+                ? "to_char(created_at, 'YYYY-MM-DD HH24:MI:SS.US')"
+                : sqlserver
+                    ? 'CONVERT(varchar(33), created_at, 126)'
+                    : 'CAST(created_at AS TEXT)';
+        let cursor;
+        for (;;) {
+            const params = provider === 'mysql' || provider === 'sqlite' ? [sessionId, sessionId] : [sessionId];
+            let before = '';
+            if (cursor) {
+                if (cursor.time === null) {
+                    before =
+                        provider === 'postgresql'
+                            ? 'AND (created_at IS NOT NULL OR (created_at IS NULL AND id < ?))'
+                            : 'AND created_at IS NULL AND id < ?';
+                    params.push(cursor.id);
+                }
+                else {
+                    before = `AND (created_at < ? OR (created_at = ? AND id < ?)${provider === 'postgresql' ? '' : ' OR created_at IS NULL'})`;
+                    params.push(cursor.time, cursor.time, cursor.id);
+                }
+            }
+            let sql = `SELECT ${sqlserver ? 'TOP 100 ' : ''}id, tool_calls, ${sortTime} AS sort_time
+        FROM ai_messages WHERE session_id = ?
+        ${provider === 'sqlite' ? 'AND session_id COLLATE BINARY = ?' : ''}
+        ${provider === 'mysql' ? 'AND BINARY session_id = BINARY ?' : ''}
+        AND deleted_at IS NULL AND history_id IS NULL AND type = 'model' AND tool_calls IS NOT NULL
+        ${sqlserver
+                ? `AND session_id COLLATE Latin1_General_100_BIN2 = @P1
+          AND DATALENGTH(CAST(session_id AS nvarchar(max))) = DATALENGTH(CAST(@P1 AS nvarchar(max)))`
+                : ''}
+        ${before} ORDER BY created_at DESC, id DESC ${sqlserver ? '' : 'LIMIT 100'}`;
+            if (sqlserver || provider === 'postgresql') {
+                let index = 0;
+                sql = sql.replaceAll('?', () => `${sqlserver ? '@P' : '$'}${++index}`);
+            }
+            const page = await this.prisma.$queryRawUnsafe(sql, ...params);
+            for (const row of page) {
+                let calls = row.tool_calls;
+                if (typeof calls === 'string') {
+                    try {
+                        calls = JSON.parse(calls);
+                    }
+                    catch {
+                        continue;
+                    }
+                }
+                if (Array.isArray(calls) && calls.some((call) => call?.id === callId)) {
+                    return this.readRows({
+                        sessionId,
+                        deletedAt: null,
+                        historyId: null,
+                        type: 'model',
+                        id: row.id,
+                    });
+                }
+            }
+            if (page.length < 100)
+                return [];
+            const last = page.at(-1);
+            cursor = { id: last.id, time: last.sort_time };
+        }
+    }
+    async readRows(where) {
         if (!this.prisma)
             return [];
-        // ponytail: tool-call lookup scans one session; add an indexed call lookup if that becomes slow.
-        // Scan only the active segment for tool results; exact-ID reads still find
-        // archived rows so redelivery does not count the same message twice.
         const rows = await this.prisma.aiMessage.findMany({
-            where: {
-                sessionId,
-                deletedAt: null,
-                ...(messageId ? { id: pk(`message ${sessionId}\0${messageId}`) } : { historyId: null }),
-            },
+            where,
             orderBy: { createdAt: 'asc' },
+            take: 1,
         });
         const json = (value) => (typeof value === 'string' ? JSON.parse(value) : value);
         return rows.map((row) => {

package/lib/index.js

@@ -20,7 +20,7 @@
 import Schema from '@deepseek-ai/schemastery';
 import { DatabaseBackend } from './backends/database.js';
 import { compressionSnapshot } from './compaction.js';
-import { mergeToolResult, projectEvent, usageSampleOf } from './projector.js';
+import { mergeToolResult, projectEvent, toolResultCallId, usageSampleOf } from './projector.js';
 export const name = 'dsh-storage';
 // Event-tap only: no hard service dependency. Add 'sessionPersistence' to
 // `inject` if you want fail-fast when persistence is not composed.
@@ -63,6 +63,54 @@ export function apply(ctx, config) {
         }),
     ];
     const sessions = new Map();
+    const toolCallMessages = new Map();
+    const TOOL_CALL_CACHE_LIMIT = 1000;
+    function rememberToolCalls(sessionId, row) {
+        if (row?.type !== 'model' || !Array.isArray(row.toolCalls))
+            return;
+        for (const call of row.toolCalls) {
+            if (typeof call?.id !== 'string')
+                continue;
+            const calls = toolCallMessages.get(sessionId) ?? new Map();
+            // Bound each session's cache; Map iteration evicts the oldest insertion first.
+            if (!calls.has(call.id) && calls.size >= TOOL_CALL_CACHE_LIMIT) {
+                const oldest = calls.keys().next().value;
+                if (oldest !== undefined)
+                    calls.delete(oldest);
+            }
+            calls.set(call.id, row.id);
+            toolCallMessages.set(sessionId, calls);
+        }
+    }
+    async function readToolResultRows(session, sessionId, event) {
+        const callId = toolResultCallId(event);
+        if (!callId)
+            return [];
+        let messageId = toolCallMessages.get(sessionId)?.get(callId);
+        if (!messageId && typeof session?.deriveMessages === 'function') {
+            const messages = session.deriveMessages();
+            for (let index = messages.length - 1; index >= 0; index--) {
+                const message = messages[index];
+                if (message.role === 'assistant' &&
+                    typeof message.id === 'string' &&
+                    Array.isArray(message.content) &&
+                    message.content.some((block) => block?.type === 'tool-call' && block.id === callId)) {
+                    messageId = message.id;
+                    break;
+                }
+            }
+        }
+        if (messageId) {
+            const rows = await backends[0].readMessages(sessionId, messageId);
+            if (rows.some((row) => Array.isArray(row.toolCalls) && row.toolCalls.some((call) => call?.id === callId)))
+                return rows;
+        }
+        // Older calls absent from the live surface are searched in bounded database pages.
+        const rows = await backends[0].readToolCallMessages(sessionId, callId);
+        for (const row of rows)
+            rememberToolCalls(sessionId, row);
+        return rows;
+    }
     // Row writes are fire-and-forget during the run, but tracked so the
     // session/flush checkpoint and shutdown can drain them — a one-shot
     // (headless) run would otherwise lose the tail events to process exit.
@@ -133,6 +181,7 @@ export function apply(ctx, config) {
         started = track(guard(async (backend) => backend.init?.()));
         return async () => {
             sessions.clear();
+            toolCallMessages.clear();
             await Promise.all([...pending]);
             await guard(async (backend) => backend.close?.());
         };
@@ -144,6 +193,7 @@ export function apply(ctx, config) {
             const archived = (chains.get(sessionId) ?? started).then(async () => {
                 const id = await backends[0].archiveSession(sessionId);
                 sessions.delete(sessionId);
+                toolCallMessages.delete(sessionId);
                 return id;
             });
             const settled = archived.then(() => { }, () => { });
@@ -178,9 +228,11 @@ export function apply(ctx, config) {
             if (row && session.header?.createdAt !== undefined) {
                 row.metadata = { ...row.metadata, sessionCreatedAt: session.header.createdAt };
             }
-            const storedRows = row || event?.type === 'tool/result'
-                ? await backends[0].readMessages(sessionId, row?.id)
-                : [];
+            const storedRows = event?.type === 'tool/result'
+                ? await readToolResultRows(session, sessionId, event)
+                : row
+                    ? await backends[0].readMessages(sessionId, row.id)
+                    : [];
             if (event?.type === 'tool/result')
                 row = mergeToolResult(event, storedRows);
             const existing = row && storedRows.find((stored) => stored.id === row?.id);
@@ -208,6 +260,7 @@ export function apply(ctx, config) {
                     accum.lastMessageAt = row.createdAt;
                 }
                 await fanout((backend) => backend.upsertMessage(row));
+                rememberToolCalls(sessionId, row);
             }
             // Replacement usage accounting: the latest sample for a (turn, step)
             // supersedes the previous one — fold only the delta, so the
@@ -263,6 +316,7 @@ export function apply(ctx, config) {
         // the dispose was meant to drop).
         enqueue(sessionId, async () => {
             sessions.delete(sessionId);
+            toolCallMessages.delete(sessionId);
         });
         // Delete the chain entry only once it has settled AND is still the same
         // entry — deleting earlier would let a recreated session race the

package/lib/projector.js

@@ -106,10 +106,14 @@ export function projectEvent(_session, event, sessionId) {
             return null;
     }
 }
+export function toolResultCallId(event) {
+    const callId = event?.data?.message?.source?.callId ?? event?.data?.message?.content?.[0]?.toolCallId;
+    return typeof callId === 'string' ? callId : undefined;
+}
 /** Update the persisted assistant call identified by a tool result. */
 export function mergeToolResult(event, rows) {
     const result = event?.data?.message?.content?.[0] ?? null;
-    const callId = event?.data?.message?.source?.callId ?? result?.toolCallId;
+    const callId = toolResultCallId(event);
     if (!callId)
         return null;
     const model = [...rows]

package/lib/types/backends/database.d.ts

@@ -31,7 +31,9 @@ export declare class DatabaseBackend implements StorageBackend {
     private jsonField;
     private toDbData;
     upsertMessage(row: MessageRow): Promise<void>;
-    readMessages(sessionId: string, messageId?: string): Promise<MessageRow[]>;
+    readMessages(sessionId: string, messageId: string): Promise<MessageRow[]>;
+    readToolCallMessages(sessionId: string, callId: string): Promise<MessageRow[]>;
+    private readRows;
     readSession(sessionId: string): Promise<SessionRow | null>;
     upsertSession(row: SessionRow): Promise<void>;
     archiveSession(sessionId: string): Promise<string | null>;

package/lib/types/projector.d.ts

@@ -22,6 +22,7 @@ import type { MessageRow } from './types.js';
  * turn/step lifecycle, approvals, ...).
  */
 export declare function projectEvent(_session: any, event: any, sessionId: string): MessageRow | null;
+export declare function toolResultCallId(event: any): string | undefined;
 /** Update the persisted assistant call identified by a tool result. */
 export declare function mergeToolResult(event: any, rows: MessageRow[]): MessageRow | null;
 export interface UsageSample {

package/lib/types/types.d.ts

@@ -43,8 +43,10 @@ export interface StorageBackend {
     init?(): Promise<void>;
     /** Read the stored session row, for seeding the live rollup on resume. */
     readSession?(sessionId: string): Promise<SessionRow | null>;
-    /** Read persisted messages, optionally narrowing to one logical message ID. */
-    readMessages(sessionId: string, messageId?: string): Promise<MessageRow[]>;
+    /** Read one persisted logical message, including archived rows for redelivery. */
+    readMessages(sessionId: string, messageId: string): Promise<MessageRow[]>;
+    /** Find at most one active model message containing this exact tool call. */
+    readToolCallMessages(sessionId: string, callId: string): Promise<MessageRow[]>;
     upsertMessage(row: MessageRow): Promise<void>;
     upsertSession(row: SessionRow): Promise<void>;
     /** Move active messages into a queryable history row; null when already empty. */

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@amaster.ai/dsh-storage",
-  "version": "0.1.33",
+  "version": "0.1.34",
   "description": "Session storage mirror for DeepSeek Harness: projects the session event stream into MySQL/PostgreSQL/SQLite/SQL Server (ai_messages / ai_chat_histories) without replacing dsh's local persistence.",
   "type": "module",
   "main": "lib/index.js",