mattrguerra/cortexmainPrivate on GitHub, published here in full

mockClient.ts

Language TypeScript19.9 KB588 lines

src/ipc/mock/mockClient.ts

// An in-memory stand-in for the Rust shell, seeded with invented chats.
// Streaming is simulated word by word; cost uses the same arithmetic as
// cortex-core's cost.rs, so the numbers shown are the numbers Rust would store.
import type { AppEvent, IpcClient } from '../client'
import type { AttachmentMeta } from '../gen/AttachmentMeta'
import type { Conversation } from '../gen/Conversation'
import type { Message } from '../gen/Message'
import type { Preset } from '../gen/Preset'
import type { Prices } from '../gen/Prices'
import type { Provider } from '../gen/Provider'
import type { Settings } from '../gen/Settings'
import type { Usage } from '../gen/Usage'
import { IMAGE_EXT, NANO } from '../../lib/format'
import { cannedReply, chats, models, PLANS_SVG, providers as seedProviders } from './fixtures'
export function costNano(u: Usage, p: Prices): number {
const cached = Math.min(Math.max(u.cached_tokens, 0), Math.max(u.prompt_tokens, 0))
const uncached = Math.max(u.prompt_tokens - cached, 0)
const cachePrice = p.cache_read_per_m ?? p.input_per_m
return Math.round(
uncached * p.input_per_m * 1000 +
cached * cachePrice * 1000 +
Math.max(u.completion_tokens, 0) * p.output_per_m * 1000,
)
}
export function estimateTokens(text: string): number {
let cjk = 0
let other = 0
for (const ch of text) {
const c = ch.codePointAt(0) ?? 0
if (
(c >= 0x3040 && c <= 0x30ff) ||
(c >= 0x3400 && c <= 0x9fff) ||
(c >= 0xac00 && c <= 0xd7af) ||
(c >= 0xf900 && c <= 0xfaff)
)
cjk++
else other++
}
return cjk + Math.floor((other + 3) / 4)
}
const defaultSettings: Settings = {
theme: 'system',
reduce_transparency: false,
material: 'sidebar',
glass_tint: 100,
shortcut: 'CommandOrControl+Shift+Space',
default_provider_id: 'deepinfra',
default_model_id: 'Qwen/Qwen3-235B-A22B-Instruct-2507',
quick_provider_id: 'deepinfra',
quick_model_id: 'google/gemma-3-12b-it',
monthly_budget_usd: 10,
confirm_over_budget: false,
vault_dir: '~/Vault/Chats',
image_links: 'wiki',
favorite_models: ['Qwen/Qwen3-235B-A22B-Instruct-2507', 'google/gemma-3-27b-it'],
vision_models: [],
send_on_enter: true,
}
let seq = 0
const id = (p: string) => `${p}-${(++seq).toString(36).padStart(4, '0')}`
export interface MockOptions {
/** Milliseconds between streamed words. */
delay?: number
empty?: boolean
hasKey?: boolean
}
export function createMockClient(opts: MockOptions = {}): IpcClient {
const delay = opts.delay ?? 22
const now = Date.now()
const providers: Provider[] = seedProviders.map((p) => ({ ...p }))
const conversations = new Map<string, Conversation>()
const messages = new Map<string, Message[]>()
const attachmentData = new Map<string, { meta: AttachmentMeta; bytes: Uint8Array }>()
const presets: Preset[] = [
{
id: 'preset-editor',
name: 'Editor',
system_prompt: "Edit for brevity. Keep the writer's voice. Prefer short sentences.",
provider_id: null,
model_id: null,
temperature: 0.3,
max_tokens: null,
},
{
id: 'preset-shell',
name: 'Shell reviewer',
system_prompt: 'You are a careful shell reviewer. Prefer small, safe changes.',
provider_id: 'deepinfra',
model_id: 'Qwen/Qwen3-235B-A22B-Instruct-2507',
temperature: null,
max_tokens: null,
},
]
let settings: Settings = { ...defaultSettings }
const keys = new Set<string>(opts.hasKey === false ? [] : ['deepinfra'])
const listeners = new Map<AppEvent, Set<(p: unknown) => void>>()
const streams = new Map<string, { cancelled: boolean }>()
const enc = new TextEncoder()
const priceOf = (provider: string, model: string): Prices | null =>
models[provider]?.find((x) => x.id === model)?.prices ?? null
const msg = (
conversation_id: string,
role: string,
content: string,
created_at: number,
extra: Partial<Message> = {},
): Message => ({
id: id('m'),
conversation_id,
role,
content,
reasoning: null,
status: 'complete',
error_kind: null,
error_message: null,
provider_id: null,
model_id: null,
usage: { prompt_tokens: 0, completion_tokens: 0, cached_tokens: 0 },
usage_source: 'none',
prices: null,
cost_nanousd: 0,
finish_reason: null,
ttft_ms: null,
duration_ms: null,
created_at,
attachments: [],
...extra,
})
if (!opts.empty) {
for (const c of chats) {
const cid = id('c')
const created = now - c.ageMs
conversations.set(cid, {
id: cid,
title: c.title,
provider_id: c.provider,
model_id: c.model,
preset_id: null,
system_prompt: c.system,
origin: 'main',
pinned: !!c.pinned,
archived: false,
created_at: created,
updated_at: created + 60_000 * c.turns.length,
})
const list: Message[] = []
c.turns.forEach((t, i) => {
const at = created + i * 60_000
const user = msg(cid, 'user', t.user, at)
if (t.image) {
const aid = id('a')
const meta: AttachmentMeta = {
id: aid,
kind: 'image',
filename: t.image.filename,
mime: 'image/svg+xml',
byte_size: PLANS_SVG.length,
width: t.image.width,
height: t.image.height,
}
attachmentData.set(aid, { meta, bytes: enc.encode(PLANS_SVG) })
user.attachments.push(meta)
}
list.push(user)
const prices = priceOf(c.provider, c.model)
const usage = { prompt_tokens: t.prompt, completion_tokens: t.completion, cached_tokens: 0 }
list.push(
msg(cid, 'assistant', t.assistant, at + 4000, {
reasoning: t.reasoning ?? null,
provider_id: c.provider,
model_id: c.model,
usage,
usage_source: t.estimated ? 'estimated' : 'reported',
prices,
cost_nanousd: prices ? costNano(usage, prices) : 0,
finish_reason: 'stop',
ttft_ms: 380 + i * 40,
duration_ms: 1800 + t.completion * 9,
}),
)
})
messages.set(cid, list)
}
}
const emit = (event: AppEvent, payload: unknown = null) => listeners.get(event)?.forEach((h) => h(payload))
const summaries = (includeArchived: boolean) =>
[...conversations.values()]
.filter((c) => includeArchived || !c.archived)
.map((c) => {
const ms = messages.get(c.id) ?? []
return {
id: c.id,
title: c.title,
model_id: c.model_id,
pinned: c.pinned,
origin: c.origin,
updated_at: c.updated_at,
cost_nanousd: ms.reduce((s, m) => s + m.cost_nanousd, 0),
cost_estimated: ms.some((m) => m.usage_source === 'estimated'),
message_count: ms.length,
has_usage: ms.some((m) => m.role === 'assistant' && m.usage_source !== 'none'),
}
})
.sort((a, b) => Number(b.pinned) - Number(a.pinned) || b.updated_at - a.updated_at)
const monthStart = () => {
const d = new Date()
return new Date(d.getFullYear(), d.getMonth(), 1).getTime()
}
const allAssistant = () => [...messages.values()].flat().filter((m) => m.role === 'assistant')
const historyTokens = (cid: string | null) => {
if (!cid) return 0
const c = conversations.get(cid)
const ms = messages.get(cid) ?? []
return (
(c?.system_prompt ? estimateTokens(c.system_prompt) + 4 : 0) +
ms
.filter((m) => m.status === 'complete')
.reduce(
(s, m) => s + estimateTokens(m.content) + 4 + m.attachments.filter((a) => a.kind === 'image').length * 765,
0,
)
)
}
const client: IpcClient = {
appInfo: async () => ({
version: '0.1.0',
platform: 'browser',
vibrancy: false,
data_dir: '(sample data, in memory)',
}),
windowReady: async () => {},
showMain: async () => {},
listProviders: async () => providers.map((p) => ({ ...p })),
saveProvider: async (p) => {
const saved = { ...p, id: p.id || id('p'), kind: p.id === 'deepinfra' ? 'deepinfra' : 'openai_compatible' }
const i = providers.findIndex((x) => x.id === saved.id)
if (i >= 0) providers[i] = saved
else providers.push(saved)
models[saved.id] ??= []
return saved
},
deleteProvider: async (pid) => {
if (pid === 'deepinfra') return
const i = providers.findIndex((x) => x.id === pid)
if (i >= 0) providers.splice(i, 1)
},
setApiKey: async (pid, key) => {
if (!key.trim()) throw new Error('The key is empty.')
if (key.trim() === 'bad-key')
throw new Error('DeepInfra rejected this key, so it was not saved. Copy the whole key again and paste it.')
keys.add(pid)
return { accepted: true, detail: null }
},
clearApiKey: async (pid) => {
keys.delete(pid)
},
hasApiKey: async (pid) => keys.has(pid),
listModels: async (pid) => ({
models: models[pid] ?? [],
source: pid === 'deepinfra' ? 'live' : 'local',
fetched_at: now - 3 * 3600_000,
snapshot_date: null,
warning: null,
}),
refreshModels: async (pid) => ({
models: models[pid] ?? [],
source: pid === 'deepinfra' ? 'live' : 'local',
fetched_at: Date.now(),
snapshot_date: null,
warning: null,
}),
listConversations: async (inc) => summaries(inc),
getConversation: async (cid) => {
const c = conversations.get(cid)
if (!c) throw new Error(`conversation ${cid} not found`)
return { conversation: { ...c }, messages: (messages.get(cid) ?? []).map((m) => ({ ...m })) }
},
createConversation: async (provider_id, model_id, preset_id, origin) => {
const preset = presets.find((p) => p.id === preset_id)
const t = Date.now()
const c: Conversation = {
id: id('c'),
title: 'New chat',
provider_id: preset?.provider_id ?? provider_id,
model_id: preset?.model_id ?? model_id,
preset_id,
system_prompt: preset?.system_prompt ?? '',
origin,
pinned: false,
archived: false,
created_at: t,
updated_at: t,
}
conversations.set(c.id, c)
messages.set(c.id, [])
return { ...c }
},
renameConversation: async (cid, title) => {
const c = conversations.get(cid)
if (c) c.title = title.trim() || 'New chat'
},
setPinned: async (cid, pinned) => {
const c = conversations.get(cid)
if (c) c.pinned = pinned
},
setArchived: async (cid, archived) => {
const c = conversations.get(cid)
if (c) c.archived = archived
},
deleteConversation: async (cid) => {
conversations.delete(cid)
messages.delete(cid)
},
setConversationModel: async (cid, provider_id, model_id) => {
const c = conversations.get(cid)
if (c) Object.assign(c, { provider_id, model_id })
},
setSystemPrompt: async (cid, system_prompt, preset_id) => {
const c = conversations.get(cid)
if (c) Object.assign(c, { system_prompt, preset_id })
},
sendMessage: async (req, onEvent) => {
const c = conversations.get(req.conversation_id)
if (!c) throw new Error('conversation not found')
const provider = providers.find((p) => p.id === c.provider_id)
if (provider?.requires_key && !keys.has(provider.id))
throw new Error(`Add your ${provider.name} API key in Settings first.`)
const list = messages.get(c.id)!
const t = Date.now()
let user: Message | null = null
if (req.text.trim() || req.attachment_ids.length) {
user = msg(c.id, 'user', req.text, t)
user.attachments = req.attachment_ids
.map((a) => attachmentData.get(a)?.meta)
.filter((a): a is AttachmentMeta => !!a)
list.push(user)
if (c.title === 'New chat')
c.title =
req.text
.split('\n')
.find((l) => l.trim())
?.trim()
.slice(0, 60) || 'New chat'
}
const lastUser = [...list].reverse().find((m) => m.role === 'user')
const images = lastUser?.attachments.filter((a) => a.kind === 'image').length ?? 0
const assistant = msg(c.id, 'assistant', '', t + 1, {
status: 'streaming',
provider_id: c.provider_id,
model_id: c.model_id,
})
list.push(assistant)
c.updated_at = t
const streamId = id('s')
const stream = { cancelled: false }
streams.set(streamId, stream)
const reply = cannedReply(lastUser?.content ?? '', images)
const words = reply.split(/(?<= )/)
const prompt = historyTokens(c.id)
;(async () => {
let text = ''
const started = Date.now()
for (const w of words) {
if (stream.cancelled) break
await new Promise((r) => setTimeout(r, delay))
if (stream.cancelled) break
text += w
onEvent({ type: 'delta', text: w })
}
const local = !provider?.requires_key
const usage = { prompt_tokens: prompt, completion_tokens: estimateTokens(text), cached_tokens: 0 }
const source = local || stream.cancelled ? 'estimated' : 'reported'
if (source === 'reported') onEvent({ type: 'usage', usage })
const prices = priceOf(c.provider_id, c.model_id)
Object.assign(assistant, {
content: text,
status: stream.cancelled ? 'cancelled' : 'complete',
usage,
usage_source: source,
prices,
cost_nanousd: prices ? costNano(usage, prices) : 0,
finish_reason: stream.cancelled ? null : 'stop',
ttft_ms: delay + 300,
duration_ms: Date.now() - started,
} satisfies Partial<Message>)
streams.delete(streamId)
onEvent({ type: 'done', message: { ...assistant } })
emit('conversations-changed')
})()
return { stream_id: streamId, user_message: user ? { ...user } : null, assistant_message: { ...assistant } }
},
cancelStream: async (sid) => {
const s = streams.get(sid)
if (s) s.cancelled = true
},
truncateFrom: async (cid, mid) => {
const list = messages.get(cid) ?? []
const i = list.findIndex((m) => m.id === mid)
if (i >= 0) list.splice(i)
},
estimatePrompt: async (a) => {
const prompt_tokens =
historyTokens(a.conversationId) +
estimateTokens(a.text) +
estimateTokens(a.extraText ?? '') +
4 +
a.imageCount * 765
const model = models[a.providerId]?.find((x) => x.id === a.modelId)
return {
prompt_tokens,
prompt_cost_nanousd: model
? costNano({ prompt_tokens, completion_tokens: 0, cached_tokens: 0 }, model.prices)
: 0,
context_length: model?.context_length ?? null,
}
},
search: async (q) => {
const words = q.toLowerCase().split(/\W+/).filter(Boolean)
if (!words.length) return []
const hits = []
for (const c of conversations.values()) {
if (c.archived) continue
for (const m of messages.get(c.id) ?? []) {
const low = m.content.toLowerCase()
if (!words.every((w) => low.includes(w))) continue
const at = low.indexOf(words[0])
const start = Math.max(0, at - 40)
const raw = m.content.slice(start, at + 80)
const marked = raw.replace(
new RegExp(`(${words.map((w) => w.replace(/[.*+?^${}()|[\]\\]/g, '\\$&')).join('|')})`, 'gi'),
'\u0001$1\u0002',
)
hits.push({
conversation_id: c.id,
title: c.title,
message_id: m.id,
snippet: (start > 0 ? '...' : '') + marked + '...',
created_at: m.created_at,
})
}
}
return hits.slice(0, 40)
},
spendSummary: async () => {
const since = monthStart()
const ms = allAssistant().filter((m) => m.created_at >= since)
const month = ms.reduce((s, m) => s + m.cost_nanousd, 0)
const budget = settings.monthly_budget_usd
const spent = month / NANO
return {
month_nanousd: month,
month_estimated: ms.some((m) => m.usage_source === 'estimated'),
budget_usd: budget,
budget_level: budget && budget > 0 ? (spent >= budget ? 2 : spent >= budget * 0.8 ? 1 : 0) : 0,
}
},
usageReport: async (days) => {
const since = days ? Date.now() - days * 86_400_000 : monthStart()
const ms = allAssistant().filter((m) => m.created_at >= since)
const group = (key: (m: Message) => string) => {
const g = new Map<
string,
{
key: string
cost_nanousd: number
prompt_tokens: number
completion_tokens: number
messages: number
estimated: boolean
}
>()
for (const m of ms) {
const k = key(m)
const r = g.get(k) ?? {
key: k,
cost_nanousd: 0,
prompt_tokens: 0,
completion_tokens: 0,
messages: 0,
estimated: false,
}
r.cost_nanousd += m.cost_nanousd
r.prompt_tokens += m.usage.prompt_tokens
r.completion_tokens += m.usage.completion_tokens
r.messages += 1
r.estimated ||= m.usage_source === 'estimated'
g.set(k, r)
}
return [...g.values()].sort((a, b) => b.cost_nanousd - a.cost_nanousd)
}
const day = (m: Message) => {
const d = new Date(m.created_at)
return `${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, '0')}-${String(d.getDate()).padStart(2, '0')}`
}
return {
by_model: group((m) => m.model_id ?? 'unknown'),
by_day: group(day).sort((a, b) => a.key.localeCompare(b.key)),
since_ms: since,
}
},
listPresets: async () => presets.map((p) => ({ ...p })),
savePreset: async (p) => {
if (!p.name.trim()) throw new Error('Give the preset a name.')
if (presets.some((x) => x.name === p.name.trim() && x.id !== p.id))
throw new Error('A preset with that name already exists.')
const saved = { ...p, id: p.id || id('preset') }
const i = presets.findIndex((x) => x.id === saved.id)
if (i >= 0) presets[i] = saved
else presets.push(saved)
return saved
},
deletePreset: async (pid) => {
const i = presets.findIndex((x) => x.id === pid)
if (i >= 0) presets.splice(i, 1)
},
getSettings: async () => ({ ...settings }),
saveSettings: async (s) => {
settings = { ...s }
emit('settings-changed', settings)
return { settings: { ...settings }, shortcut_error: null }
},
pickAttachments: async () => [],
attachBytes: async (filename, bytes) => {
const isImage = IMAGE_EXT.test(filename)
const meta: AttachmentMeta = {
id: id('a'),
kind: isImage ? 'image' : 'text',
filename,
mime: isImage ? 'image/png' : 'text/plain',
byte_size: bytes.length,
width: isImage ? 800 : null,
height: isImage ? 600 : null,
}
attachmentData.set(meta.id, { meta, bytes })
return meta
},
attachmentBytes: async (aid) => {
const b = attachmentData.get(aid)?.bytes ?? new Uint8Array()
return b.buffer.slice(b.byteOffset, b.byteOffset + b.byteLength) as ArrayBuffer
},
discardAttachment: async (aid) => {
attachmentData.delete(aid)
},
exportConversation: async (cid) => ({
path: `${settings.vault_dir ?? '~/Vault'}/${conversations.get(cid)?.title ?? 'Chat'}.md`,
files_copied: 0,
}),
chooseVaultFolder: async () => settings.vault_dir,
quickOpenInMain: async () => {},
hideQuick: async () => {},
resizeQuick: async () => {},
openUrl: async (url) => {
window.open(url, '_blank', 'noopener')
},
on: async (event, handler) => {
const set = listeners.get(event) ?? new Set()
set.add(handler)
listeners.set(event, set)
return () => set.delete(handler)
},
}
return client
}