@fickydev/pigent 0.1.5 → 0.1.7

package/src/db/repositories/HeartbeatRepository.ts

@@ -0,0 +1,129 @@
+import { and, desc, eq, gte, isNull } from "drizzle-orm";
+import { nanoid } from "nanoid";
+import type { DbClient } from "../client";
+import { heartbeats, type HeartbeatRow } from "../schema";
+
+export type HeartbeatStatus = "pending" | "running" | "noop" | "notified" | "failed";
+
+export type CreateHeartbeatInput = {
+  agentId: string;
+  sessionId?: string | null;
+  prompt: string;
+  status?: Extract<HeartbeatStatus, "pending" | "running">;
+  startedAt?: number | null;
+};
+
+export type FinishHeartbeatInput = {
+  result?: string | null;
+  error?: string | null;
+  finishedAt?: number;
+};
+
+export class HeartbeatRepository {
+  constructor(private readonly db: DbClient) {}
+
+  async create(input: CreateHeartbeatInput): Promise<HeartbeatRow> {
+    const id = nanoid();
+    const now = Date.now();
+
+    await this.db.insert(heartbeats).values({
+      id,
+      agentId: input.agentId,
+      sessionId: input.sessionId ?? null,
+      status: input.status ?? "pending",
+      prompt: input.prompt,
+      result: null,
+      error: null,
+      startedAt: input.startedAt ?? null,
+      finishedAt: null,
+      createdAt: now,
+    });
+
+    const row = await this.findById(id);
+    if (!row) throw new Error(`failed to create heartbeat ${id}`);
+
+    return row;
+  }
+
+  async markRunning(id: string, startedAt = Date.now()): Promise<HeartbeatRow> {
+    await this.db
+      .update(heartbeats)
+      .set({
+        status: "running",
+        startedAt,
+      })
+      .where(eq(heartbeats.id, id));
+
+    return this.requireById(id);
+  }
+
+  async markNoop(id: string, input: FinishHeartbeatInput = {}): Promise<HeartbeatRow> {
+    return this.finish(id, "noop", input);
+  }
+
+  async markNotified(id: string, input: FinishHeartbeatInput = {}): Promise<HeartbeatRow> {
+    return this.finish(id, "notified", input);
+  }
+
+  async markFailed(id: string, input: FinishHeartbeatInput): Promise<HeartbeatRow> {
+    return this.finish(id, "failed", input);
+  }
+
+  async findById(id: string): Promise<HeartbeatRow | null> {
+    const row = await this.db.query.heartbeats.findFirst({
+      where: eq(heartbeats.id, id),
+    });
+
+    return row ?? null;
+  }
+
+  async findLatest(agentId: string, sessionId?: string | null): Promise<HeartbeatRow | null> {
+    const sessionPredicate = sessionId ? eq(heartbeats.sessionId, sessionId) : isNull(heartbeats.sessionId);
+
+    const row = await this.db.query.heartbeats.findFirst({
+      where: and(eq(heartbeats.agentId, agentId), sessionPredicate),
+      orderBy: desc(heartbeats.createdAt),
+    });
+
+    return row ?? null;
+  }
+
+  async countRecentNotified(agentId: string, since: number, sessionId?: string | null): Promise<number> {
+    const sessionPredicate = sessionId ? eq(heartbeats.sessionId, sessionId) : isNull(heartbeats.sessionId);
+
+    const rows = await this.db.query.heartbeats.findMany({
+      where: and(
+        eq(heartbeats.agentId, agentId),
+        sessionPredicate,
+        eq(heartbeats.status, "notified"),
+        gte(heartbeats.finishedAt, since),
+      ),
+      columns: {
+        id: true,
+      },
+    });
+
+    return rows.length;
+  }
+
+  private async finish(id: string, status: Extract<HeartbeatStatus, "noop" | "notified" | "failed">, input: FinishHeartbeatInput): Promise<HeartbeatRow> {
+    await this.db
+      .update(heartbeats)
+      .set({
+        status,
+        result: input.result ?? null,
+        error: input.error ?? null,
+        finishedAt: input.finishedAt ?? Date.now(),
+      })
+      .where(eq(heartbeats.id, id));
+
+    return this.requireById(id);
+  }
+
+  private async requireById(id: string): Promise<HeartbeatRow> {
+    const row = await this.findById(id);
+    if (!row) throw new Error(`heartbeat not found ${id}`);
+
+    return row;
+  }
+}

package/src/db/repositories/SessionRepository.ts

@@ -3,6 +3,8 @@ import { nanoid } from "nanoid";
 import type { DbClient } from "../client";
 import { agentSessions, type AgentSessionRow } from "../schema";
 
+export type SessionThinkingLevel = "off" | "low" | "medium" | "high";
+
 export type SessionKey = {
   agentId: string;
   channel: string;
@@ -38,6 +40,30 @@ export class SessionRepository {
     return created;
   }
 
+  async updateModel(id: string, model: string | null): Promise<AgentSessionRow> {
+    await this.db
+      .update(agentSessions)
+      .set({
+        model,
+        updatedAt: Date.now(),
+      })
+      .where(eq(agentSessions.id, id));
+
+    return this.requireById(id);
+  }
+
+  async updateThinkingLevel(id: string, thinkingLevel: SessionThinkingLevel | null): Promise<AgentSessionRow> {
+    await this.db
+      .update(agentSessions)
+      .set({
+        thinkingLevel,
+        updatedAt: Date.now(),
+      })
+      .where(eq(agentSessions.id, id));
+
+    return this.requireById(id);
+  }
+
   async findById(id: string): Promise<AgentSessionRow | null> {
     const row = await this.db.query.agentSessions.findFirst({
       where: eq(agentSessions.id, id),
@@ -46,6 +72,13 @@ export class SessionRepository {
     return row ?? null;
   }
 
+  private async requireById(id: string): Promise<AgentSessionRow> {
+    const row = await this.findById(id);
+    if (!row) throw new Error(`session not found ${id}`);
+
+    return row;
+  }
+
   private async findByKey(key: SessionKey): Promise<AgentSessionRow | null> {
     const threadPredicate = key.threadId
       ? eq(agentSessions.threadId, key.threadId)

package/src/db/repositories/index.ts

@@ -1,5 +1,6 @@
 import { db } from "../client";
 import { AgentRepository } from "./AgentRepository";
+import { HeartbeatRepository } from "./HeartbeatRepository";
 import { MessageRepository } from "./MessageRepository";
 import { RuntimeKvRepository } from "./RuntimeKvRepository";
 import { SessionRepository } from "./SessionRepository";
@@ -10,6 +11,7 @@ export type Repositories = ReturnType<typeof createRepositories>;
 export function createRepositories(dbClient = db) {
   return {
     agents: new AgentRepository(dbClient),
+    heartbeats: new HeartbeatRepository(dbClient),
     messages: new MessageRepository(dbClient),
     runtimeKv: new RuntimeKvRepository(dbClient),
     sessions: new SessionRepository(dbClient),

package/src/db/schema.ts

@@ -48,6 +48,8 @@ export const agentSessions = sqliteTable(
     userId: text("user_id"),
     piSessionId: text("pi_session_id"),
     instructionsHash: text("instructions_hash"),
+    model: text("model"),
+    thinkingLevel: text("thinking_level", { enum: ["off", "low", "medium", "high"] }),
     createdAt: integer("created_at").notNull(),
     updatedAt: integer("updated_at").notNull(),
   },

package/src/pi/PiAgentRunner.ts

@@ -10,10 +10,13 @@ import {
 import { mkdir } from "node:fs/promises";
 import { resolve } from "node:path";
 import type { LoadedAgentConfig, ProfileConfig } from "../config/schemas";
+import type { AgentSessionRow } from "../db/schema";
+import { resolveModelSelection } from "./PiModelResolver";
 
 export type PiAgentRunInput = {
   agent: LoadedAgentConfig;
   profile: ProfileConfig | null;
+  session: AgentSessionRow;
   prompt: string;
 };
 
@@ -37,10 +40,13 @@ export class PiAgentRunner {
 
     const authStorage = AuthStorage.create();
     const modelRegistry = ModelRegistry.create(authStorage);
+    const modelSelection = resolveModelSelection(modelRegistry, [input.session, input.agent, input.profile ?? {}]);
     const { session } = await createAgentSession({
       cwd: workspace,
       authStorage,
       modelRegistry,
+      model: modelSelection.model,
+      thinkingLevel: modelSelection.thinkingLevel,
       resourceLoader,
       sessionManager: SessionManager.create(workspace),
       settingsManager,

package/src/pi/PiModelResolver.ts

@@ -0,0 +1,76 @@
+import type { ThinkingLevel } from "@earendil-works/pi-agent-core";
+import type { Model } from "@earendil-works/pi-ai";
+import type { ModelRegistry } from "@earendil-works/pi-coding-agent";
+import { logger } from "../logging/logger";
+
+export type ModelSelectionConfig = {
+  model?: string | null;
+  thinkingLevel?: ThinkingLevel | null;
+};
+
+export type ResolvedModelSelection = {
+  model?: Model<any>;
+  thinkingLevel?: ThinkingLevel;
+};
+
+export function resolveModelSelection(
+  modelRegistry: ModelRegistry,
+  configs: ModelSelectionConfig[],
+): ResolvedModelSelection {
+  const configuredModel = firstConfigured(configs.map((config) => config.model));
+  const thinkingLevel = firstConfigured(configs.map((config) => config.thinkingLevel));
+  const resolved: ResolvedModelSelection = {};
+
+  if (configuredModel) {
+    const model = resolveModel(modelRegistry, configuredModel);
+    if (model) {
+      resolved.model = model;
+    }
+  }
+
+  if (thinkingLevel) {
+    resolved.thinkingLevel = thinkingLevel;
+  }
+
+  return resolved;
+}
+
+function resolveModel(modelRegistry: ModelRegistry, modelRef: string): Model<any> | undefined {
+  const parsed = parseModelRef(modelRef);
+
+  if (!parsed) {
+    logger.warn("invalid model reference; expected provider/modelId", { model: modelRef });
+    return undefined;
+  }
+
+  const model = modelRegistry.find(parsed.provider, parsed.modelId);
+
+  if (!model) {
+    logger.warn("configured model not found; falling back to Pi default", { model: modelRef });
+    return undefined;
+  }
+
+  return model;
+}
+
+export function isValidModelRef(modelRef: string): boolean {
+  return parseModelRef(modelRef) !== null;
+}
+
+function parseModelRef(modelRef: string): { provider: string; modelId: string } | null {
+  const trimmed = modelRef.trim();
+  const separatorIndex = trimmed.indexOf("/");
+
+  if (separatorIndex <= 0 || separatorIndex === trimmed.length - 1) {
+    return null;
+  }
+
+  return {
+    provider: trimmed.slice(0, separatorIndex),
+    modelId: trimmed.slice(separatorIndex + 1),
+  };
+}
+
+function firstConfigured<T>(values: Array<T | null | undefined>): T | undefined {
+  return values.find((value): value is T => value !== null && value !== undefined && value !== "");
+}