348 lines
11 KiB
TypeScript
348 lines
11 KiB
TypeScript
import { randomUUID } from 'crypto';
|
|
import { join } from 'path';
|
|
import type { PiEvent, MessageCost, Message } from '@@/api/pi/types';
|
|
import { sessionManager } from '@@/api/pi/session-manager';
|
|
import * as storage from '@@/api/pi/storage';
|
|
import * as piBridge from '@@/api/pi/pi-bridge';
|
|
import { getHomeDirForRole } from '@@/data-path';
|
|
import { getUserSettings } from 'officerdb';
|
|
import { logger } from '@@/api/pi/logger';
|
|
import { sendClaudeCode, clearClaudeCodeSession } from './send-claude-code';
|
|
|
|
const DEFAULT_MODEL = 'opencode/big-pickle';
|
|
const IDLE_TIMEOUT_MS = 60 * 60 * 1000;
|
|
const SEND_TIMEOUT_MS = 5 * 60 * 1000;
|
|
|
|
type SendAndAwaitParams = {
|
|
userId: number;
|
|
email: string;
|
|
username: string;
|
|
prompt: string;
|
|
context: string;
|
|
contextId: string;
|
|
model?: string;
|
|
role?: string;
|
|
};
|
|
|
|
type SendAndAwaitResult = {
|
|
text: string;
|
|
sessionId: string;
|
|
model: string;
|
|
cost: MessageCost;
|
|
};
|
|
|
|
// Per-session mutex to serialize concurrent prompts
|
|
const sessionLocks = new Map<string, Promise<void>>();
|
|
|
|
// Per-session callback — swapped each time a new prompt is sent
|
|
type EventCallback = (event: PiEvent) => void;
|
|
const sessionCallbacks = new Map<string, EventCallback>();
|
|
|
|
// Channel model overrides — survive session eviction/recreation
|
|
const channelModelOverrides = new Map<string, string>();
|
|
|
|
function buildSessionId(context: string, userId: number, contextId: string): string {
|
|
return `channel-${context}-${userId}-${contextId}`;
|
|
}
|
|
|
|
async function getUserDefaultModel(userId: number): Promise<string | null> {
|
|
try {
|
|
const settings = await getUserSettings(userId);
|
|
const chat = settings?.chat as Record<string, unknown> | undefined;
|
|
return (chat?.defaultModel as string) || null;
|
|
} catch {
|
|
return null;
|
|
}
|
|
}
|
|
|
|
export function getSessionModel(context: string, userId: number, contextId: string): string | null {
|
|
const sessionId = buildSessionId(context, userId, contextId);
|
|
const session = sessionManager.getSession(sessionId);
|
|
return session?.model ?? channelModelOverrides.get(sessionId) ?? null;
|
|
}
|
|
|
|
export function setSessionModel(context: string, userId: number, contextId: string, model: string): void {
|
|
const sessionId = buildSessionId(context, userId, contextId);
|
|
// Store override independently of session — survives idle eviction
|
|
channelModelOverrides.set(sessionId, model);
|
|
clearClaudeCodeSession(sessionId);
|
|
|
|
const session = sessionManager.getSession(sessionId);
|
|
if (session) {
|
|
session.model = model;
|
|
session.meta.model = model;
|
|
// Kill existing PI process so it respawns with the new model
|
|
if (session.piProcess) {
|
|
piBridge.killPi(session.piProcess);
|
|
session.piProcess = null;
|
|
}
|
|
logger.info('Channel model switched', { sessionId, model, killedProcess: true });
|
|
} else {
|
|
logger.info('Channel model override stored (no active session)', { sessionId, model });
|
|
}
|
|
}
|
|
|
|
export async function sendAndAwait(params: SendAndAwaitParams): Promise<SendAndAwaitResult> {
|
|
const { userId, context, contextId } = params;
|
|
const sessionId = buildSessionId(context, userId, contextId);
|
|
|
|
// Serialize per session — if two messages arrive at once, second waits for first
|
|
const existing = sessionLocks.get(sessionId) ?? Promise.resolve();
|
|
let releaseLock: () => void;
|
|
const lockPromise = new Promise<void>((resolve) => {
|
|
releaseLock = resolve;
|
|
});
|
|
sessionLocks.set(
|
|
sessionId,
|
|
existing.then(() => lockPromise),
|
|
);
|
|
|
|
await existing;
|
|
|
|
try {
|
|
// Resolve model early to check for claude-code routing
|
|
const override = channelModelOverrides.get(sessionId);
|
|
const resolvedModel = params.model ?? override ?? (await getUserDefaultModel(params.userId)) ?? DEFAULT_MODEL;
|
|
|
|
if (resolvedModel.startsWith('claude-code')) {
|
|
return await sendClaudeCode({
|
|
userId: params.userId,
|
|
email: params.email,
|
|
username: params.username,
|
|
prompt: params.prompt,
|
|
sessionKey: sessionId,
|
|
model: resolvedModel,
|
|
});
|
|
}
|
|
|
|
return await doSend(sessionId, params);
|
|
} finally {
|
|
releaseLock!();
|
|
if (sessionLocks.get(sessionId) === existing.then(() => lockPromise)) {
|
|
sessionLocks.delete(sessionId);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Persistent event dispatcher — registered once at spawn time, delegates to current callback
|
|
function createDispatcher(sessionId: string): (event: PiEvent) => void {
|
|
return (event: PiEvent) => {
|
|
const cb = sessionCallbacks.get(sessionId);
|
|
if (cb) cb(event);
|
|
};
|
|
}
|
|
|
|
async function doSend(sessionId: string, params: SendAndAwaitParams): Promise<SendAndAwaitResult> {
|
|
const { userId, email, username, prompt, context, contextId } = params;
|
|
const homeDir = getHomeDirForRole(email, params.role ?? null);
|
|
const cwd = homeDir;
|
|
|
|
// Resolve model: explicit param > !model override > user default > existing session > system default
|
|
const existingSession = sessionManager.getSession(sessionId);
|
|
const override = channelModelOverrides.get(sessionId);
|
|
let model = params.model ?? override;
|
|
if (!model) {
|
|
const userDefault = await getUserDefaultModel(userId);
|
|
model = userDefault ?? existingSession?.model ?? DEFAULT_MODEL;
|
|
}
|
|
|
|
const session = sessionManager.getOrCreate(sessionId, email, cwd, model, null, context, contextId);
|
|
session.model = model;
|
|
session.meta.model = model;
|
|
session.userId = userId;
|
|
logger.info('Channel doSend', { sessionId, model, hasProcess: !!session.piProcess, userId, email });
|
|
|
|
return new Promise<SendAndAwaitResult>((resolve, reject) => {
|
|
let resultText = '';
|
|
const cost: MessageCost = { inputTokens: 0, outputTokens: 0, totalUSD: 0 };
|
|
let settled = false;
|
|
|
|
const timeout = setTimeout(() => {
|
|
if (!settled) {
|
|
settled = true;
|
|
session.isGenerating = false;
|
|
sessionCallbacks.delete(sessionId);
|
|
reject(new Error('sendAndAwait timed out after 5 minutes'));
|
|
}
|
|
}, SEND_TIMEOUT_MS);
|
|
|
|
const settle = () => {
|
|
sessionCallbacks.delete(sessionId);
|
|
clearTimeout(timeout);
|
|
settled = true;
|
|
};
|
|
|
|
// Register per-prompt callback — the persistent dispatcher will call this
|
|
sessionCallbacks.set(sessionId, (event: PiEvent) => {
|
|
if (settled) return;
|
|
|
|
switch (event.type) {
|
|
case 'delta': {
|
|
session.streamBuffer += event.text;
|
|
break;
|
|
}
|
|
|
|
case 'text': {
|
|
const text = event.text || session.streamBuffer;
|
|
if (text) {
|
|
resultText += (resultText ? '\n\n' : '') + text;
|
|
const assistantMsg: Message = {
|
|
id: randomUUID(),
|
|
timestamp: Date.now(),
|
|
role: 'assistant',
|
|
text,
|
|
model,
|
|
};
|
|
session.messages.push(assistantMsg);
|
|
session.meta.messageCount += 1;
|
|
session.streamBuffer = '';
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'tool:start': {
|
|
if (session.streamBuffer) {
|
|
resultText += (resultText ? '\n\n' : '') + session.streamBuffer;
|
|
const assistantMsg: Message = {
|
|
id: randomUUID(),
|
|
timestamp: Date.now(),
|
|
role: 'assistant',
|
|
text: session.streamBuffer,
|
|
model,
|
|
};
|
|
session.messages.push(assistantMsg);
|
|
session.meta.messageCount += 1;
|
|
session.streamBuffer = '';
|
|
}
|
|
const toolMsg: Message = {
|
|
id: randomUUID(),
|
|
timestamp: Date.now(),
|
|
role: 'tool',
|
|
toolCallId: event.toolCallId,
|
|
toolName: event.toolName,
|
|
toolInput: event.toolInput,
|
|
};
|
|
session.messages.push(toolMsg);
|
|
session.meta.messageCount += 1;
|
|
break;
|
|
}
|
|
|
|
case 'tool:result': {
|
|
for (let i = session.messages.length - 1; i >= 0; i--) {
|
|
const m = session.messages[i]!;
|
|
if (m.role === 'tool' && m.toolCallId === event.toolCallId) {
|
|
m.output = event.output;
|
|
m.isError = event.isError;
|
|
break;
|
|
}
|
|
}
|
|
break;
|
|
}
|
|
|
|
case 'result': {
|
|
if (session.streamBuffer) {
|
|
resultText += (resultText ? '\n\n' : '') + session.streamBuffer;
|
|
const assistantMsg: Message = {
|
|
id: randomUUID(),
|
|
timestamp: Date.now(),
|
|
role: 'assistant',
|
|
text: session.streamBuffer,
|
|
model,
|
|
cost: event.cost,
|
|
};
|
|
session.messages.push(assistantMsg);
|
|
session.meta.messageCount += 1;
|
|
session.streamBuffer = '';
|
|
}
|
|
|
|
session.isGenerating = false;
|
|
session.meta.cost.inputTokens += event.cost.inputTokens;
|
|
session.meta.cost.outputTokens += event.cost.outputTokens;
|
|
session.meta.cost.totalUSD += event.cost.totalUSD;
|
|
session.meta.updatedAt = Date.now();
|
|
|
|
cost.inputTokens = event.cost.inputTokens;
|
|
cost.outputTokens = event.cost.outputTokens;
|
|
cost.totalUSD = event.cost.totalUSD;
|
|
|
|
storage.saveSession(homeDir, sessionId, session.meta, session.messages).catch((err) => {
|
|
logger.error('Failed to save channel session', { sessionId, error: String(err) });
|
|
});
|
|
|
|
sessionManager.setIdleTimeout(sessionId, IDLE_TIMEOUT_MS);
|
|
|
|
settle();
|
|
resolve({ text: resultText || '(no response)', sessionId, model: model!, cost });
|
|
break;
|
|
}
|
|
|
|
case 'error': {
|
|
session.isGenerating = false;
|
|
settle();
|
|
reject(new Error(event.message));
|
|
break;
|
|
}
|
|
|
|
case 'stopped': {
|
|
session.isGenerating = false;
|
|
settle();
|
|
resolve({ text: resultText || '(stopped)', sessionId, model: model!, cost });
|
|
break;
|
|
}
|
|
}
|
|
});
|
|
|
|
// Spawn Pi process if not running
|
|
(async () => {
|
|
try {
|
|
if (!session.piProcess) {
|
|
const role = params.role;
|
|
let spawnOptions: { sessionFile?: string; username?: string; role?: string } | undefined;
|
|
if (session.messages.length > 0) {
|
|
await storage.saveSession(homeDir, sessionId, session.meta, session.messages);
|
|
const hostPath = await storage.getSessionFilePath(homeDir, sessionId);
|
|
if (hostPath) {
|
|
spawnOptions = { sessionFile: hostPath, username, role };
|
|
}
|
|
}
|
|
if (!spawnOptions) spawnOptions = { username, role };
|
|
|
|
const dispatcher = createDispatcher(sessionId);
|
|
session.piProcess = await piBridge.spawnPi(cwd, model!, userId, email, dispatcher, spawnOptions);
|
|
|
|
const proc = session.piProcess;
|
|
proc.exited.then(() => {
|
|
if (session.piProcess === proc) {
|
|
session.piProcess = null;
|
|
logger.info('Channel Pi process exited', { sessionId });
|
|
}
|
|
});
|
|
|
|
logger.info('Spawned Pi for channel session', { sessionId, model, context, userId, email });
|
|
}
|
|
|
|
// Add user message
|
|
const userMsg: Message = {
|
|
id: randomUUID(),
|
|
timestamp: Date.now(),
|
|
role: 'user',
|
|
text: prompt,
|
|
};
|
|
session.messages.push(userMsg);
|
|
session.meta.messageCount += 1;
|
|
session.meta.updatedAt = Date.now();
|
|
|
|
if (!session.meta.title) {
|
|
session.meta.title = prompt.slice(0, 100);
|
|
}
|
|
|
|
session.isGenerating = true;
|
|
piBridge.sendPrompt(session.piProcess, prompt, randomUUID());
|
|
} catch (err) {
|
|
settle();
|
|
reject(err);
|
|
}
|
|
})();
|
|
});
|
|
}
|