William Valentin
334 lines
12 KiB
TypeScript
334 lines
12 KiB
TypeScript
import { randomUUID } from 'crypto';
|
|
import type { SessionManager } from '../session/manager.js';
|
|
import type { ModelClient } from '../models/types.js';
|
|
import type { ModelRouter, ModelTier } from '../models/router.js';
|
|
import type { ToolRegistry } from '../tools/registry.js';
|
|
import type { ToolExecutor } from '../tools/executor.js';
|
|
import { AgentOrchestrator, type DelegationConfig } from '../backends/native/orchestrator.js';
|
|
import type { ToolUseEvent } from '../backends/native/agent.js';
|
|
import type { MemoryStore } from '../memory/store.js';
|
|
import type { Config } from '../config/index.js';
|
|
import { summarizeSessionOnEnd, type SessionEndSummaryConfig } from '../session/endSummary.js';
|
|
|
|
export interface SessionBridgeConfig {
|
|
sessionManager: SessionManager;
|
|
modelClient: ModelClient | ModelRouter;
|
|
systemPrompt: string;
|
|
toolRegistry: ToolRegistry;
|
|
toolExecutor: ToolExecutor;
|
|
config?: Config;
|
|
memoryStore?: MemoryStore;
|
|
}
|
|
|
|
interface ClientEntry {
|
|
connectionId: string;
|
|
sessionId: string;
|
|
agent: AgentOrchestrator;
|
|
busy: boolean;
|
|
}
|
|
|
|
export class SessionBridge {
|
|
private clients: Map<string, ClientEntry> = new Map();
|
|
private agents: Map<string, AgentOrchestrator> = new Map();
|
|
private config: SessionBridgeConfig;
|
|
|
|
constructor(config: SessionBridgeConfig) {
|
|
this.config = config;
|
|
}
|
|
|
|
/** Register a new WS connection. Returns the assigned connection ID. */
|
|
connect(connectionId?: string): string {
|
|
const id = connectionId ?? randomUUID();
|
|
const sessionId = `ws:${id}`;
|
|
const agent = this.getOrCreateAgent(sessionId);
|
|
|
|
this.clients.set(id, {
|
|
connectionId: id,
|
|
sessionId,
|
|
agent,
|
|
busy: false,
|
|
});
|
|
|
|
return id;
|
|
}
|
|
|
|
/** Remove a WS connection. Does NOT destroy the session (persists in SQLite). */
|
|
disconnect(connectionId: string): void {
|
|
const client = this.clients.get(connectionId);
|
|
if (client) {
|
|
// Only remove the agent if no other clients share the session
|
|
const otherClients = Array.from(this.clients.values())
|
|
.filter(c => c.sessionId === client.sessionId && c.connectionId !== connectionId);
|
|
if (otherClients.length === 0) {
|
|
const agent = this.agents.get(client.sessionId);
|
|
const summaryConfig = this.config.config?.sessions?.end_summary;
|
|
if (agent && summaryConfig?.enabled) {
|
|
const history = agent.getHistory();
|
|
const mappedConfig: SessionEndSummaryConfig = {
|
|
enabled: summaryConfig.enabled,
|
|
tier: summaryConfig.tier,
|
|
maxMessages: summaryConfig.max_messages,
|
|
maxInputChars: summaryConfig.max_input_chars,
|
|
maxTokens: summaryConfig.max_tokens,
|
|
writeToMemory: summaryConfig.write_to_memory,
|
|
memoryNamespace: summaryConfig.memory_namespace,
|
|
};
|
|
void summarizeSessionOnEnd({
|
|
agent,
|
|
sessionId: client.sessionId,
|
|
history,
|
|
config: mappedConfig,
|
|
memoryStore: this.config.memoryStore,
|
|
}).catch((error) => {
|
|
console.warn('Session end summary failed:', error);
|
|
});
|
|
}
|
|
this.agents.delete(client.sessionId);
|
|
}
|
|
this.clients.delete(connectionId);
|
|
}
|
|
}
|
|
|
|
/** Switch a connection to a different session (e.g. resuming an old session). */
|
|
switchSession(connectionId: string, sessionId: string): void {
|
|
const client = this.clients.get(connectionId);
|
|
if (!client) {throw new Error(`Unknown connection: ${connectionId}`);}
|
|
if (client.busy) {throw new Error('Cannot switch session while agent is busy');}
|
|
|
|
const agent = this.getOrCreateAgent(sessionId);
|
|
client.sessionId = sessionId;
|
|
client.agent = agent;
|
|
}
|
|
|
|
/** Get the AgentOrchestrator for a connection. */
|
|
getAgent(connectionId: string): AgentOrchestrator | undefined {
|
|
return this.clients.get(connectionId)?.agent;
|
|
}
|
|
|
|
/** Get the session ID for a connection. */
|
|
getSessionId(connectionId: string): string | undefined {
|
|
return this.clients.get(connectionId)?.sessionId;
|
|
}
|
|
|
|
/** Check if a connection's agent is busy. */
|
|
isBusy(connectionId: string): boolean {
|
|
return this.clients.get(connectionId)?.busy ?? false;
|
|
}
|
|
|
|
/** Mark a connection's agent as busy/idle. */
|
|
setBusy(connectionId: string, busy: boolean): void {
|
|
const client = this.clients.get(connectionId);
|
|
if (client) {client.busy = busy;}
|
|
}
|
|
|
|
/** Request cancellation for the current operation on a connection's agent. */
|
|
cancel(connectionId: string): boolean {
|
|
const client = this.clients.get(connectionId);
|
|
if (!client || !client.busy) {
|
|
return false;
|
|
}
|
|
|
|
client.agent.cancel();
|
|
return true;
|
|
}
|
|
|
|
/**
|
|
* Request cancellation for the active operation in a session.
|
|
* Returns true if at least one connection in the session is currently busy.
|
|
*/
|
|
cancelSession(sessionId: string): boolean {
|
|
const clients = Array.from(this.clients.values()).filter((client) => client.sessionId === sessionId);
|
|
if (clients.length === 0) {
|
|
return false;
|
|
}
|
|
|
|
const hasBusyClient = clients.some((client) => client.busy);
|
|
if (!hasBusyClient) {
|
|
return false;
|
|
}
|
|
|
|
const agent = this.agents.get(sessionId);
|
|
if (!agent) {
|
|
return false;
|
|
}
|
|
|
|
agent.cancel();
|
|
return true;
|
|
}
|
|
|
|
/** Set onToolUse callback for a connection's agent. */
|
|
setOnToolUse(connectionId: string, callback: ((event: ToolUseEvent) => void) | undefined): void {
|
|
const client = this.clients.get(connectionId);
|
|
if (client) {client.agent.setOnToolUse(callback);}
|
|
}
|
|
|
|
/** List all active sessions with connection counts. */
|
|
listSessions(): Array<{ sessionId: string; connections: number }> {
|
|
const sessionMap = new Map<string, number>();
|
|
for (const client of this.clients.values()) {
|
|
sessionMap.set(client.sessionId, (sessionMap.get(client.sessionId) ?? 0) + 1);
|
|
}
|
|
return Array.from(sessionMap.entries()).map(([sessionId, connections]) => ({
|
|
sessionId,
|
|
connections,
|
|
}));
|
|
}
|
|
|
|
/** Get count of active connections. */
|
|
get connectionCount(): number {
|
|
return this.clients.size;
|
|
}
|
|
|
|
/** Get usage stats for a specific connection's agent. */
|
|
getUsage(connectionId: string): { inputTokens: number; outputTokens: number; calls: number } | undefined {
|
|
const agent = this.clients.get(connectionId)?.agent;
|
|
if (!agent) {return undefined;}
|
|
const usage = agent.getUsage();
|
|
return {
|
|
inputTokens: usage.primary.inputTokens,
|
|
outputTokens: usage.primary.outputTokens,
|
|
calls: usage.primary.calls,
|
|
};
|
|
}
|
|
|
|
/** Get usage stats for all active sessions. Returns an array of per-session usage entries. */
|
|
getAllUsage(): Array<{
|
|
sessionId: string;
|
|
primary: { inputTokens: number; outputTokens: number; calls: number };
|
|
delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
|
|
total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
|
|
}> {
|
|
const results: Array<{
|
|
sessionId: string;
|
|
primary: { inputTokens: number; outputTokens: number; calls: number };
|
|
delegation: Record<string, { inputTokens: number; outputTokens: number; calls: number }>;
|
|
total: { inputTokens: number; outputTokens: number; calls: number; estimatedCost: number };
|
|
}> = [];
|
|
|
|
// De-duplicate by sessionId (multiple connections may share a session)
|
|
const seen = new Set<string>();
|
|
|
|
for (const client of this.clients.values()) {
|
|
if (seen.has(client.sessionId)) {continue;}
|
|
seen.add(client.sessionId);
|
|
|
|
const usage = client.agent.getUsage();
|
|
results.push({
|
|
sessionId: client.sessionId,
|
|
primary: usage.primary,
|
|
delegation: usage.delegation,
|
|
total: usage.total,
|
|
});
|
|
}
|
|
|
|
return results;
|
|
}
|
|
|
|
/** Get estimated context budget for all active sessions. */
|
|
getAllContextUsage(): Array<{
|
|
sessionId: string;
|
|
budget: {
|
|
estimatedTokens: number;
|
|
contextWindow: number;
|
|
remainingTokens: number;
|
|
usagePct: number;
|
|
thresholdPct: number;
|
|
thresholdTokens: number;
|
|
shouldCompact: boolean;
|
|
};
|
|
}> {
|
|
const results: Array<{
|
|
sessionId: string;
|
|
budget: {
|
|
estimatedTokens: number;
|
|
contextWindow: number;
|
|
remainingTokens: number;
|
|
usagePct: number;
|
|
thresholdPct: number;
|
|
thresholdTokens: number;
|
|
shouldCompact: boolean;
|
|
};
|
|
}> = [];
|
|
|
|
const seen = new Set<string>();
|
|
for (const client of this.clients.values()) {
|
|
if (seen.has(client.sessionId)) {
|
|
continue;
|
|
}
|
|
seen.add(client.sessionId);
|
|
results.push({
|
|
sessionId: client.sessionId,
|
|
budget: client.agent.getContextBudget(),
|
|
});
|
|
}
|
|
|
|
return results;
|
|
}
|
|
|
|
private getOrCreateAgent(sessionId: string): AgentOrchestrator {
|
|
let agent = this.agents.get(sessionId);
|
|
if (!agent) {
|
|
const session = this.config.sessionManager.getSession('ws', sessionId);
|
|
const config = this.config.config;
|
|
|
|
// Read per-session tier override from session config
|
|
const sessionTier = session.getConfig?.('modelTier') as ModelTier | undefined;
|
|
const primaryTier = sessionTier ?? config?.agents.primary_tier ?? 'default';
|
|
|
|
const delegationConfig: DelegationConfig = {
|
|
compaction: config?.agents.delegation.compaction ?? 'fast',
|
|
memory_extraction: config?.agents.delegation.memory_extraction ?? 'fast',
|
|
classification: config?.agents.delegation.classification ?? 'fast',
|
|
tool_summarisation: config?.agents.delegation.tool_summarisation ?? 'fast',
|
|
complex_reasoning: config?.agents.delegation.complex_reasoning ?? 'complex',
|
|
};
|
|
|
|
agent = new AgentOrchestrator({
|
|
modelRouter: this.config.modelClient as ModelRouter,
|
|
systemPrompt: this.config.systemPrompt,
|
|
session,
|
|
toolRegistry: this.config.toolRegistry,
|
|
toolExecutor: this.config.toolExecutor,
|
|
primaryTier,
|
|
delegation: delegationConfig,
|
|
maxDelegationDepth: config?.agents.max_delegation_depth ?? 3,
|
|
maxIterations: config?.agents.max_iterations,
|
|
compaction: config?.compaction.enabled ? {
|
|
thresholdPct: config.compaction.threshold_pct,
|
|
keepTurns: config.compaction.keep_turns,
|
|
summaryMaxTokens: config.compaction.summary_max_tokens,
|
|
importanceThreshold: config.compaction.importance_threshold,
|
|
proactive: {
|
|
enabled: config.compaction.proactive.enabled,
|
|
warnPct: config.compaction.proactive.warn_pct,
|
|
checkpointPct: config.compaction.proactive.checkpoint_pct,
|
|
autoCompactPct: config.compaction.proactive.auto_compact_pct,
|
|
checkpointCooldownMs: config.compaction.proactive.checkpoint_cooldown_ms,
|
|
memoryNamespace: config.compaction.proactive.memory_namespace,
|
|
},
|
|
} : undefined,
|
|
modelName: config?.models.default.model,
|
|
contextWindow: config?.models.default.context_window,
|
|
memoryStore: this.config.memoryStore,
|
|
memoryAutoExtract: config?.memory?.auto_extract,
|
|
memoryInjectionStrategy: config?.memory?.injection_strategy,
|
|
memoryMaxInjectionTokens: config?.memory?.max_injection_tokens,
|
|
toolPolicyContext: {
|
|
agent: primaryTier,
|
|
provider: config?.models.default.provider,
|
|
autonomyLevel: config?.agents.autonomy_level ?? 'standard',
|
|
sensitiveMode: config?.agents.sensitive_mode ?? 'deny_without_elevation',
|
|
immutableDenylist: (config?.agents.immutable_denylist ?? []).map((rule) => ({
|
|
tool: rule.tool,
|
|
argsPattern: rule.args_pattern,
|
|
reason: rule.reason,
|
|
})),
|
|
},
|
|
});
|
|
|
|
this.agents.set(sessionId, agent);
|
|
}
|
|
return agent;
|
|
}
|
|
}
|