@rudderhq/server 0.1.0-canary.9 → 0.1.0

package/dist/services/runtime-kernel/heartbeat.js

@@ -1,13 +1,14 @@
 import path from "node:path";
 import { and, asc, desc, eq, gt, gte, inArray, lte, sql } from "drizzle-orm";
-import { agents, agentRuntimeState, agentTaskSessions, agentWakeupRequests, activityLog, heartbeatRunEvents, heartbeatRuns, issueComments, issues, projects, } from "@rudderhq/db";
+import { AGENT_RUN_CONCURRENCY_DEFAULT, AGENT_RUN_CONCURRENCY_MAX, AGENT_RUN_CONCURRENCY_MIN, } from "@rudderhq/shared";
+import { agents, agentRuntimeState, agentTaskSessions, agentWakeupRequests, activityLog, heartbeatRunEvents, heartbeatRuns, issueComments, issues, organizations, projects, } from "@rudderhq/db";
 import { conflict, notFound } from "../../errors.js";
 import { createExecutionScores, observeExecutionEvent, updateExecutionObservation, updateExecutionTraceIO, updateExecutionTraceName, updateExecutionTraceSession, withExecutionObservation, } from "../../langfuse.js";
 import { emitExecutionTranscriptTree } from "../../langfuse-transcript.js";
 import { logger } from "../../middleware/logger.js";
 import { publishLiveEvent } from "../live-events.js";
 import { getRunLogStore } from "../run-log-store.js";
-import { getServerAdapter, runningProcesses } from "../../agent-runtimes/index.js";
+import { findServerAdapter, getServerAdapter, runningProcesses } from "../../agent-runtimes/index.js";
 import { createLocalAgentJwt } from "../../agent-auth-jwt.js";
 import { parseObject, asBoolean, asNumber, appendWithCap, MAX_EXCERPT_BYTES } from "../../agent-runtimes/utils.js";
 import { costService } from "../costs.js";
@@ -27,10 +28,12 @@ import { logActivity } from "../activity-log.js";
 import { redactCurrentUserText, redactCurrentUserValue } from "../../log-redaction.js";
 import { hasSessionCompactionThresholds, resolveSessionCompactionPolicy, } from "@rudderhq/agent-runtime-utils";
 import { buildCreateAgentBenchmarkTags, coerceCreateAgentBenchmarkMetadata, extractCreateAgentBenchmarkMetadata, } from "@rudderhq/run-intelligence-core";
+import { executeAdapterWithModelFallbacks } from "./model-fallback.js";
 export { prioritizeProjectWorkspaceCandidatesForRun } from "../agent-run-context.js";
 const MAX_LIVE_LOG_CHUNK_BYTES = 8 * 1024;
-const HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT = 1;
-const HEARTBEAT_MAX_CONCURRENT_RUNS_MAX = 10;
+const HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT = AGENT_RUN_CONCURRENCY_DEFAULT;
+const HEARTBEAT_MAX_CONCURRENT_RUNS_MIN = AGENT_RUN_CONCURRENCY_MIN;
+const HEARTBEAT_MAX_CONCURRENT_RUNS_MAX = AGENT_RUN_CONCURRENCY_MAX;
 const DEFERRED_WAKE_CONTEXT_KEY = "_paperclipWakeContext";
 const DETACHED_PROCESS_ERROR_CODE = "process_detached";
 const ORPHANED_PROCESS_TERMINATION_GRACE_MS = 2_000;
@@ -118,7 +121,7 @@ function normalizeMaxConcurrentRuns(value) {
     const parsed = Math.floor(asNumber(value, HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT));
     if (!Number.isFinite(parsed))
         return HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT;
-    return Math.max(HEARTBEAT_MAX_CONCURRENT_RUNS_DEFAULT, Math.min(HEARTBEAT_MAX_CONCURRENT_RUNS_MAX, parsed));
+    return Math.max(HEARTBEAT_MAX_CONCURRENT_RUNS_MIN, Math.min(HEARTBEAT_MAX_CONCURRENT_RUNS_MAX, parsed));
 }
 async function withAgentStartLock(agentId, fn) {
     const previous = startLocksByAgent.get(agentId) ?? Promise.resolve();
@@ -173,9 +176,20 @@ export function buildHeartbeatRuntimeTraceMetadata(input) {
     };
 }
 export function buildHeartbeatAdapterInvokePayload(input) {
+    const explicitUsedSkills = Array.isArray(input.meta.usedSkills)
+        ? input.meta.usedSkills
+            .map((entry) => normalizeLoadedSkill(entry))
+            .filter((entry) => Boolean(entry))
+        : [];
+    const usedSkills = explicitUsedSkills.length > 0
+        ? explicitUsedSkills
+        : inferUsedSkillsFromPrompt(input.meta.prompt, input.runtimeSkills);
     return {
         ...input.meta,
         ...summarizeRuntimeSkillsForTrace(input.runtimeSkills),
+        usedSkillCount: usedSkills.length,
+        usedSkillKeys: usedSkills.map((entry) => entry.key),
+        usedSkills,
     };
 }
 function buildRecentDateKeys(windowDays, now) {
@@ -221,6 +235,92 @@ function normalizeLoadedSkill(value) {
     const label = rawRuntimeName ?? rawName ?? fallbackSkillLabel(key);
     return { key, label };
 }
+function normalizeSkillCandidate(value) {
+    return value
+        ?.trim()
+        .replace(/^\$/u, "")
+        .replace(/[?#].*$/u, "")
+        .replace(/\/+$/u, "")
+        .toLowerCase() || "";
+}
+function addSkillCandidate(candidates, value) {
+    const normalized = normalizeSkillCandidate(value);
+    if (!normalized)
+        return;
+    candidates.add(normalized);
+    const lastSegment = normalized.split(/[/:]/u).filter(Boolean).at(-1);
+    if (lastSegment)
+        candidates.add(lastSegment);
+}
+function readSkillReferenceSlug(href) {
+    const normalized = href.trim().replace(/[?#].*$/u, "").replace(/\/+$/u, "");
+    if (!normalized)
+        return null;
+    if (normalized.endsWith("/SKILL.md")) {
+        return normalized.slice(0, -"/SKILL.md".length).split("/").filter(Boolean).at(-1) ?? null;
+    }
+    if (normalized.toLowerCase().endsWith(".md")) {
+        const fileName = normalized.split("/").filter(Boolean).at(-1) ?? "";
+        return fileName.replace(/\.md$/iu, "") || null;
+    }
+    return null;
+}
+function collectSkillReferences(prompt) {
+    const references = [];
+    const pattern = /\[([^\]\n]+)\]\(([^)\n]+(?:\/SKILL\.md|\.md)(?:[?#][^)\n]*)?)\)/giu;
+    for (const match of prompt.matchAll(pattern)) {
+        const rawLabel = match[1]?.trim() ?? "";
+        const href = match[2]?.trim() ?? "";
+        if (!rawLabel || !href)
+            continue;
+        const labelWithoutPrefix = rawLabel.replace(/^\$/u, "").trim();
+        const slug = readSkillReferenceSlug(href);
+        const isExplicitSkillToken = rawLabel.startsWith("$") || href.replace(/[?#].*$/u, "").endsWith("/SKILL.md");
+        if (!isExplicitSkillToken)
+            continue;
+        const key = labelWithoutPrefix || slug;
+        if (!key)
+            continue;
+        const label = slug ?? fallbackSkillLabel(key);
+        const candidates = new Set();
+        addSkillCandidate(candidates, labelWithoutPrefix);
+        addSkillCandidate(candidates, slug);
+        addSkillCandidate(candidates, href);
+        references.push({ key, label, candidates });
+    }
+    return references;
+}
+function inferUsedSkillsFromPrompt(prompt, loadedSkills) {
+    const promptText = readNonEmptyString(prompt);
+    if (!promptText)
+        return [];
+    const references = collectSkillReferences(promptText);
+    if (references.length === 0)
+        return [];
+    const loaded = loadedSkills
+        .map((entry) => normalizeLoadedSkill(entry))
+        .filter((entry) => Boolean(entry))
+        .map((entry) => {
+        const candidates = new Set();
+        addSkillCandidate(candidates, entry.key);
+        addSkillCandidate(candidates, entry.label);
+        return { ...entry, candidates };
+    });
+    const used = new Map();
+    for (const reference of references) {
+        const matched = loaded.find((entry) => {
+            for (const candidate of reference.candidates) {
+                if (entry.candidates.has(candidate))
+                    return true;
+            }
+            return false;
+        });
+        const normalized = matched ?? { key: reference.key, label: reference.label };
+        if (!used.has(normalized.key))
+            used.set(normalized.key, normalized);
+    }
+    return Array.from(used.values());
+}
 function normalizeLedgerBillingType(value) {
     const raw = readNonEmptyString(value);
     switch (raw) {
@@ -1201,6 +1301,7 @@ export function heartbeatService(db) {
         const existing = await getRuntimeState(agent.id);
         if (existing)
             return existing;
+        const now = new Date();
         return db
             .insert(agentRuntimeState)
             .values({
@@ -1208,6 +1309,16 @@ export function heartbeatService(db) {
             orgId: agent.orgId,
             agentRuntimeType: agent.agentRuntimeType,
             stateJson: {},
+            createdAt: now,
+            updatedAt: now,
+        })
+            .onConflictDoUpdate({
+            target: agentRuntimeState.agentId,
+            set: {
+                orgId: agent.orgId,
+                agentRuntimeType: agent.agentRuntimeType,
+                updatedAt: now,
+            },
         })
             .returning()
             .then((rows) => rows[0]);
@@ -2183,6 +2294,7 @@ export function heartbeatService(db) {
                 let stderrTranscriptBuffer = "";
                 let stdoutTranscriptParser = null;
                 let transcriptFallbackResult = null;
+                let modelTurnInput;
                 let finalObservationOutput = null;
                 let finalObservationStatus = run.status;
                 let finalObservationSessionId = heartbeatObservationContext.sessionKey ?? null;
@@ -2701,6 +2813,7 @@ export function heartbeatService(db) {
                                     meta.env[key] = "***REDACTED***";
                             }
                         }
+                        modelTurnInput = meta.prompt;
                         heartbeatObservationContext.metadata = {
                             ...(heartbeatObservationContext.metadata ?? {}),
                             ...buildHeartbeatRuntimeTraceMetadata({
@@ -2734,7 +2847,7 @@ export function heartbeatService(db) {
                             agentRuntimeType: agent.agentRuntimeType,
                         }, "local agent jwt secret missing or invalid; running without injected RUDDER_API_KEY");
                     }
-                    const adapterResult = await adapter.execute({
+                    const adapterResult = await executeAdapterWithModelFallbacks(adapter, {
                         runId: run.id,
                         agent,
                         runtime: runtimeForAdapter,
@@ -2746,6 +2859,12 @@ export function heartbeatService(db) {
                             await persistRunProcessMetadata(run.id, meta);
                         },
                         authToken: authToken ?? undefined,
+                    }, {
+                        resolveAdapter: findServerAdapter,
+                        createAuthToken: (agentRuntimeType) => createLocalAgentJwt(agent.id, agent.orgId, agentRuntimeType, run.id) ?? undefined,
+                        onAttemptStart: (_attempt, attemptAdapter) => {
+                            stdoutTranscriptParser = attemptAdapter.parseStdoutLine ?? null;
+                        },
                     });
                     const adapterManagedRuntimeServices = adapterResult.runtimeServices
                         ? await persistAdapterManagedRuntimeServices({
@@ -3028,6 +3147,7 @@ export function heartbeatService(db) {
                             context: heartbeatObservationContext,
                             parentObservation: observation,
                             transcript: executionTranscript,
+                            initialTurnInput: modelTurnInput,
                             fallbackResult: transcriptFallbackResult,
                         });
                         finalObservationOutput = transcriptStats.finalOutput ?? transcriptFallbackResult?.output ?? null;
@@ -4092,6 +4212,104 @@ export function heartbeatService(db) {
             now: opts?.now ?? new Date(),
         });
     }
+    async function buildSkillAnalytics(scope, opts) {
+        const now = opts?.now ?? new Date();
+        const customDateKeys = opts?.startDate && opts?.endDate
+            ? buildDateKeysBetween(opts.startDate, opts.endDate).slice(0, 120)
+            : [];
+        const windowDays = customDateKeys.length > 0
+            ? customDateKeys.length
+            : Math.max(1, Math.min(opts?.windowDays ?? 30, 90));
+        const dateKeys = customDateKeys.length > 0
+            ? customDateKeys
+            : buildRecentDateKeys(windowDays, now);
+        const startDate = dateKeys[0];
+        const endDate = dateKeys.at(-1);
+        const windowStart = new Date(`${startDate}T00:00:00.000Z`);
+        const windowEnd = new Date(`${endDate}T23:59:59.999Z`);
+        const rows = await db
+            .select({
+            createdAt: heartbeatRunEvents.createdAt,
+            payload: heartbeatRunEvents.payload,
+        })
+            .from(heartbeatRunEvents)
+            .where(and(eq(heartbeatRunEvents.orgId, scope.orgId), ...(scope.agentId ? [eq(heartbeatRunEvents.agentId, scope.agentId)] : []), eq(heartbeatRunEvents.eventType, "adapter.invoke"), gte(heartbeatRunEvents.createdAt, windowStart), lte(heartbeatRunEvents.createdAt, windowEnd)))
+            .orderBy(asc(heartbeatRunEvents.createdAt), asc(heartbeatRunEvents.id));
+        const days = new Map();
+        for (const date of dateKeys) {
+            days.set(date, { totalCount: 0, runCount: 0, skills: new Map() });
+        }
+        const overallSkills = new Map();
+        let totalCount = 0;
+        let totalRunsWithSkills = 0;
+        for (const row of rows) {
+            const date = new Date(row.createdAt).toISOString().slice(0, 10);
+            const bucket = days.get(date);
+            if (!bucket)
+                continue;
+            const payload = parseObject(row.payload);
+            const loadedSkills = Array.isArray(payload.loadedSkills) ? payload.loadedSkills : [];
+            const usedSkills = Array.isArray(payload.usedSkills)
+                ? payload.usedSkills
+                : inferUsedSkillsFromPrompt(payload.prompt, loadedSkills);
+            if (usedSkills.length === 0)
+                continue;
+            const eventSkills = new Map();
+            for (const entry of usedSkills) {
+                const normalized = normalizeLoadedSkill(entry);
+                if (!normalized)
+                    continue;
+                if (!eventSkills.has(normalized.key)) {
+                    eventSkills.set(normalized.key, normalized.label);
+                }
+            }
+            if (eventSkills.size === 0)
+                continue;
+            bucket.runCount += 1;
+            totalRunsWithSkills += 1;
+            for (const [key, label] of eventSkills) {
+                bucket.totalCount += 1;
+                totalCount += 1;
+                const existingDaySkill = bucket.skills.get(key);
+                if (existingDaySkill) {
+                    existingDaySkill.count += 1;
+                }
+                else {
+                    bucket.skills.set(key, { key, label, count: 1 });
+                }
+                const existingOverallSkill = overallSkills.get(key);
+                if (existingOverallSkill) {
+                    existingOverallSkill.count += 1;
+                }
+                else {
+                    overallSkills.set(key, { key, label, count: 1 });
+                }
+            }
+        }
+        return {
+            agentId: scope.agentId ?? "__all__",
+            orgId: scope.orgId,
+            windowDays,
+            startDate,
+            endDate,
+            totalCount,
+            totalRunsWithSkills,
+            skills: Array.from(overallSkills.values()).sort((left, right) => (right.count - left.count
+                || left.label.localeCompare(right.label)
+                || left.key.localeCompare(right.key))),
+            days: dateKeys.map((date) => {
+                const bucket = days.get(date);
+                return {
+                    date,
+                    totalCount: bucket.totalCount,
+                    runCount: bucket.runCount,
+                    skills: Array.from(bucket.skills.values()).sort((left, right) => (right.count - left.count
+                        || left.label.localeCompare(right.label)
+                        || left.key.localeCompare(right.key))),
+                };
+            }),
+        };
+    }
     return {
         list: async (orgId, agentId, limit) => {
             const query = db
@@ -4111,99 +4329,18 @@ export function heartbeatService(db) {
             const agent = await getAgent(agentId);
             if (!agent)
                 throw notFound("Agent not found");
-            const now = opts?.now ?? new Date();
-            const customDateKeys = opts?.startDate && opts?.endDate
-                ? buildDateKeysBetween(opts.startDate, opts.endDate).slice(0, 120)
-                : [];
-            const windowDays = customDateKeys.length > 0
-                ? customDateKeys.length
-                : Math.max(1, Math.min(opts?.windowDays ?? 30, 90));
-            const dateKeys = customDateKeys.length > 0
-                ? customDateKeys
-                : buildRecentDateKeys(windowDays, now);
-            const startDate = dateKeys[0];
-            const endDate = dateKeys.at(-1);
-            const windowStart = new Date(`${startDate}T00:00:00.000Z`);
-            const windowEnd = new Date(`${endDate}T23:59:59.999Z`);
-            const rows = await db
-                .select({
-                createdAt: heartbeatRunEvents.createdAt,
-                payload: heartbeatRunEvents.payload,
-            })
-                .from(heartbeatRunEvents)
-                .where(and(eq(heartbeatRunEvents.orgId, agent.orgId), eq(heartbeatRunEvents.agentId, agent.id), eq(heartbeatRunEvents.eventType, "adapter.invoke"), gte(heartbeatRunEvents.createdAt, windowStart), lte(heartbeatRunEvents.createdAt, windowEnd)))
-                .orderBy(asc(heartbeatRunEvents.createdAt), asc(heartbeatRunEvents.id));
-            const days = new Map();
-            for (const date of dateKeys) {
-                days.set(date, { totalCount: 0, runCount: 0, skills: new Map() });
-            }
-            const overallSkills = new Map();
-            let totalCount = 0;
-            let totalRunsWithSkills = 0;
-            for (const row of rows) {
-                const date = new Date(row.createdAt).toISOString().slice(0, 10);
-                const bucket = days.get(date);
-                if (!bucket)
-                    continue;
-                const payload = parseObject(row.payload);
-                const loadedSkills = Array.isArray(payload.loadedSkills) ? payload.loadedSkills : [];
-                if (loadedSkills.length === 0)
-                    continue;
-                const eventSkills = new Map();
-                for (const entry of loadedSkills) {
-                    const normalized = normalizeLoadedSkill(entry);
-                    if (!normalized)
-                        continue;
-                    if (!eventSkills.has(normalized.key)) {
-                        eventSkills.set(normalized.key, normalized.label);
-                    }
-                }
-                if (eventSkills.size === 0)
-                    continue;
-                bucket.runCount += 1;
-                totalRunsWithSkills += 1;
-                for (const [key, label] of eventSkills) {
-                    bucket.totalCount += 1;
-                    totalCount += 1;
-                    const existingDaySkill = bucket.skills.get(key);
-                    if (existingDaySkill) {
-                        existingDaySkill.count += 1;
-                    }
-                    else {
-                        bucket.skills.set(key, { key, label, count: 1 });
-                    }
-                    const existingOverallSkill = overallSkills.get(key);
-                    if (existingOverallSkill) {
-                        existingOverallSkill.count += 1;
-                    }
-                    else {
-                        overallSkills.set(key, { key, label, count: 1 });
-                    }
-                }
-            }
-            return {
-                agentId: agent.id,
-                orgId: agent.orgId,
-                windowDays,
-                startDate,
-                endDate,
-                totalCount,
-                totalRunsWithSkills,
-                skills: Array.from(overallSkills.values()).sort((left, right) => (right.count - left.count
-                    || left.label.localeCompare(right.label)
-                    || left.key.localeCompare(right.key))),
-                days: dateKeys.map((date) => {
-                    const bucket = days.get(date);
-                    return {
-                        date,
-                        totalCount: bucket.totalCount,
-                        runCount: bucket.runCount,
-                        skills: Array.from(bucket.skills.values()).sort((left, right) => (right.count - left.count
-                            || left.label.localeCompare(right.label)
-                            || left.key.localeCompare(right.key))),
-                    };
-                }),
-            };
+            return buildSkillAnalytics({ orgId: agent.orgId, agentId: agent.id }, opts);
+        },
+        getOrganizationSkillAnalytics: async (orgId, opts) => {
+            const org = await db
+                .select({ id: organizations.id })
+                .from(organizations)
+                .where(eq(organizations.id, orgId))
+                .limit(1)
+                .then((rows) => rows[0] ?? null);
+            if (!org)
+                throw notFound("Organization not found");
+            return buildSkillAnalytics({ orgId }, opts);
         },
         getRun,
         getRuntimeState: async (agentId) => {