diff --git a/src/main/agent-runtime/runtime-bridge.ts b/src/main/agent-runtime/runtime-bridge.ts index 169d3f3a..ac9f2299 100644 --- a/src/main/agent-runtime/runtime-bridge.ts +++ b/src/main/agent-runtime/runtime-bridge.ts @@ -5,6 +5,8 @@ import { ensureTaskThread } from "@main/thread-integration/thread-helpers"; import { SqliteSessionStore } from "./sqlite-session-store"; import { HumanApprovalHook } from "./approval-hook"; import { CostWriter } from "./cost-writer"; +import type { CredentialResolver } from "@main/credentials"; +import { getConnectorDefinition } from "@main/connectors"; export interface RuntimeBridgeDeps { db: AppDatabase; @@ -14,6 +16,7 @@ export interface RuntimeBridgeDeps { warn: (...args: unknown[]) => void; error: (...args: unknown[]) => void; }; + credentialResolver: CredentialResolver; } interface ActiveRunContext { @@ -155,7 +158,7 @@ export class RuntimeBridge { /** Handles RPC requests issued by the worker (`rpc-request` envelope). */ async handleRpc(request: { rpcId: string; - method: "beforeToolCall" | "afterToolCall" | "persistMessage"; + method: "beforeToolCall" | "afterToolCall" | "persistMessage" | "prepareRun"; runId: string; payload: unknown; }): Promise<{ ok: true; result: unknown } | { ok: false; error: string }> { @@ -184,6 +187,37 @@ export class RuntimeBridge { this.getOrCreateStore(payload.taskSessionId).appendMessage(payload.message); return { ok: true, result: undefined }; } + if (request.method === "prepareRun") { + const payload = request.payload as { + connectorId: string; + taskSessionId: string; + agentId: string; + companyId: string; + taskId: string | null; + modelId: string; + }; + this.registerActiveRun({ + runId: request.runId, + agentId: payload.agentId, + companyId: payload.companyId, + taskId: payload.taskId, + connectorId: payload.connectorId, + modelId: payload.modelId, + taskSessionId: payload.taskSessionId, + }); + const definition = getConnectorDefinition(payload.connectorId); + const resolved = await this.deps.credentialResolver.resolve(payload.connectorId); + const priorMessages = this.getOrCreateStore(payload.taskSessionId).loadMessages(); + return { + ok: true, + result: { + provider: definition.provider, + apiKey: resolved.apiKey, + extraEnv: resolved.extraEnv ?? {}, + priorMessages, + }, + }; + } return { ok: false, error: `Unknown RPC method: ${(request as { method: string }).method}` }; } catch (error) { return { ok: false, error: error instanceof Error ? error.message : String(error) }; diff --git a/src/main/orchestration/turn-context.ts b/src/main/orchestration/turn-context.ts index 30a87293..783ae3f8 100644 --- a/src/main/orchestration/turn-context.ts +++ b/src/main/orchestration/turn-context.ts @@ -133,18 +133,25 @@ function resolveSession( deps: Pick, ) { const activeSession = db.getActiveTaskSession(task.id); - let taskSession = activeSession + let taskSession: { + id: string | null; + sessionParams: Record | null; + sessionDisplayId: string | null; + } = activeSession ? { - sessionParams: activeSession.sessionParamsJson ? JSON.parse(activeSession.sessionParamsJson) as Record : null, + id: activeSession.id, + sessionParams: activeSession.sessionParamsJson + ? (JSON.parse(activeSession.sessionParamsJson) as Record) + : null, sessionDisplayId: activeSession.sessionDisplayId, } - : db.getTaskSession(task.id); + : { id: null, ...db.getTaskSession(task.id) }; let finalHandoff = sessionHandoff; if (deps.evaluatePreRunSessionRotation && activeSession) { const rotationResult = deps.evaluatePreRunSessionRotation(task.id, connectorId); if (rotationResult.rotated) { - taskSession = { sessionParams: null, sessionDisplayId: null }; + taskSession = { id: null, sessionParams: null, sessionDisplayId: null }; if (!finalHandoff && rotationResult.handoff) { finalHandoff = rotationResult.handoff; } @@ -220,7 +227,7 @@ function buildConnectorCatalog( ) { return connectors .filter((c) => c.status === "ready") - .filter((c) => c.id.startsWith("claude_compat") || db.hasModelPoolSettings(c.id)) + .filter((c) => db.hasModelPoolSettings(c.id)) .map((c) => ({ id: c.id, label: c.label, @@ -336,5 +343,7 @@ export function buildRunPayload(params: RunDispatchParams, deps: TurnContextDeps systemPrompt, trigger, managedSkills: skills.map((s) => ({ runtimeName: s.runtimeName, sourcePath: s.sourcePath })), + taskSessionId: taskSession.id ?? runId, + executionMode: db.getCompanyExecutionMode(companyId) as "ai_ceo" | "ai_cofounder", }; } diff --git a/src/main/orchestration/types.ts b/src/main/orchestration/types.ts index 577fc52b..e08995bd 100644 --- a/src/main/orchestration/types.ts +++ b/src/main/orchestration/types.ts @@ -44,6 +44,10 @@ export interface EnqueueRunPayload { systemPrompt: string | null; trigger: string | null; managedSkills: Array<{ runtimeName: string; sourcePath: string }>; + /** Task session id used by SqliteSessionStore to load/append messages. */ + taskSessionId: string; + /** Execution mode from the owning company. */ + executionMode: "ai_ceo" | "ai_cofounder"; } export interface OrchestrationDependencies { diff --git a/src/main/service-graph.ts b/src/main/service-graph.ts index a9f249b1..104c1692 100644 --- a/src/main/service-graph.ts +++ b/src/main/service-graph.ts @@ -14,6 +14,7 @@ import { createWorkerService } from "./worker-service"; import { createAppLifecycleService } from "./app-lifecycle-service"; import { createBudgetEnforcementService } from "./budget-enforcement-service"; import { createSessionManagementService } from "./session-management-service"; +import { createPiRuntimeServices, type PiRuntimeServices } from "./service-graph-pi-runtime"; import type { DesktopEvent, RunRecord } from "@shared/types"; import type { AutomationEventContext } from "./automation-service"; import type { NotificationOptions } from "./notification-service"; @@ -82,6 +83,7 @@ export function createServiceGraph({ let skillService: ReturnType | null = null; let budgetEnforcementService: ReturnType | null = null; let sessionManagementService: ReturnType | null = null; + const piRuntimeRef: { current: PiRuntimeServices | null } = { current: null }; function getOrCreateSkillService() { if (!skillService) throw new Error("Skill service is not initialized."); @@ -184,10 +186,22 @@ export function createServiceGraph({ }); desktopShellService.initialize(); + piRuntimeRef.current = createPiRuntimeServices({ + db, + emitEvent, + logger, + refreshConnectorHealth: async (connectorId: string) => { + if (connectorHealthService) { + await connectorHealthService.checkConnector(connectorId as never); + } + }, + }); + connectorHealthService = createConnectorHealthService({ db, logger, publishDomainChanged, + credentialStore: piRuntimeRef.current.credentialStore, }); workerService = createWorkerService({ @@ -205,6 +219,7 @@ export function createServiceGraph({ notify, budgetService: getOrCreateBudgetEnforcementService(), sessionService: getOrCreateSessionManagementService(), + getRuntimeBridge: () => piRuntimeRef.current?.runtimeBridge ?? null, }); orchestrationService = createOrchestrationService({ @@ -324,6 +339,8 @@ export function createServiceGraph({ skillService = null; budgetEnforcementService = null; sessionManagementService = null; + piRuntimeRef.current?.release(); + piRuntimeRef.current = null; db.close(); // Force exit after grace period — covers all platforms (macOS Dock, Windows tray, Linux) // Worker child processes (claude/codex CLI) may keep the event loop alive indefinitely @@ -346,5 +363,6 @@ export function createServiceGraph({ getOrCreateSkillService, getOrCreateBudgetEnforcementService, getOrCreateSessionManagementService, + getPiRuntime: () => piRuntimeRef.current, }; } diff --git a/src/main/session-management-service.ts b/src/main/session-management-service.ts index 7790c6b4..fe80ddb5 100644 --- a/src/main/session-management-service.ts +++ b/src/main/session-management-service.ts @@ -1,6 +1,5 @@ import type { AppDatabase } from "./database"; -import { getConnectorDefinition } from "./connectors"; -import type { ConnectorId, TaskSessionRecord } from "@shared/types"; +import type { TaskSessionRecord } from "@shared/types"; interface SessionManagementDependencies { db: AppDatabase; @@ -38,20 +37,11 @@ export function createSessionManagementService({ db, logger }: SessionManagement // company-level overrides; omitted fields fall back to DEFAULT_THRESHOLDS. function evaluateRotation( session: TaskSessionRecord, - connectorId: string, + _connectorId: string, thresholdsOverride?: RotationThresholds, ): { shouldRotate: boolean; reason: string | null } { - // D-04: Check adapter exemption - try { - const connector = getConnectorDefinition(connectorId as ConnectorId); - if (connector.module.nativeContextManagement === true) { - return { shouldRotate: false, reason: null }; - } - } catch { - // Unknown connector - proceed with rotation check - } - - // Merge overrides with defaults (per D-01 configurability) + // Pi-mono delegates context management to pi-agent-core compaction hooks; + // AgentCompany always applies its own rotation thresholds. const thresholds = { ...DEFAULT_THRESHOLDS, ...thresholdsOverride }; // D-01: Any-of thresholds - first trigger wins