Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
132 kB · 3741 lines
TypeScript
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742import type { AgentMessage } from "@earendil-works/pi-agent-core";import { getModel } from "@earendil-works/pi-ai";import { describe, expect, it, vi } from "vitest";import { createManualContinuationMessage, MANUAL_CONTINUATION_PROMPT } from "../src/core/messages.js";import { MissingSessionCwdError } from "../src/core/session-cwd.js";import { SessionImportFileNotFoundError } from "../src/core/session-import-errors.js";import { DAEMON_REFINE_REQUEST_TIMEOUT_MS, DaemonAgentConnection,} from "../src/modes/agent-connection/daemon-agent-connection.js";import type { AgentConnectionEvent, AgentConnectionRlmChildAgentSnapshot, AgentConnectionSavedSessionInfo, AgentConnectionState,} from "../src/modes/agent-connection/types.js";
import { DaemonCapabilityUnavailableError, type DaemonClientCloseListener, type DaemonClientMessageListener, type DaemonClientRequestOptions, type DaemonHello, DaemonSocketClosedError, type DaemonTransportClient,} from "../src/modes/daemon/daemon-client.js";import { DAEMON_PROTOCOL_INFO, DAEMON_SCHEMA_REVISION, type DaemonAttachResult, type DaemonCommand, type DaemonOutbound, type DaemonResponse,} from "../src/modes/daemon/daemon-protocol.js";import { DaemonRoutedClient } from "../src/modes/daemon/daemon-routed-client.js";import type { DaemonWorkerClient } from "../src/modes/daemon/daemon-worker-client.js";
class FakeDaemonClient { readonly requests: DaemonCommand[] = []; readonly requestTimeouts: number[] = []; attachResultFactory: ((command: Extract<DaemonCommand, { type: "attach" }>) => DaemonAttachResult) | undefined; restoredAttachGate: Promise<void> | undefined; restoredAttachCompleted = 0; closeCount = 0; emitCloseOnClose = false; connected = true; reconnectCount = 0; resetTransportCount = 0; reconnectError: Error | undefined; attachFailures = 0; connectionStateGate: Promise<void> | undefined; connectionStateFactory: ((activeSessionId: string) => AgentConnectionState) | undefined; rlmChildren: AgentConnectionRlmChildAgentSnapshot[] = []; rlmChildrenEventSequence = 12; rlmChildrenGate: Promise<void> | undefined; abortBashUnknownCommand = false; abortAndClearQueueUnknownCommand = false; inputPauseAcquireGate: Promise<void> | undefined; cronAddGate: Promise<void> | undefined; promptGate: Promise<void> | undefined; promptError: Error | undefined; promptResponseError: string | undefined; cancelPromptAdmissionStatus: "cancelled" | "owned" | "unknown" = "owned"; serverCapabilities = new Set<string>(); updateRestartSessions: Array<Record<string, unknown>> = []; hello: DaemonHello | undefined = { type: "daemon_hello", socketPath: "/tmp/fake.sock", protocol: DAEMON_PROTOCOL_INFO, schemaRevision: DAEMON_SCHEMA_REVISION, clientId: "fake-client", serverCapabilities: ["prompt_admission_cancellation", "session_input_admission"], }; private readonly messageListeners = new Set<DaemonClientMessageListener>(); private readonly closeListeners = new Set<DaemonClientCloseListener>();
async request( command: DaemonCommand, timeoutMs = 30000, options: DaemonClientRequestOptions = {}, ): Promise<DaemonResponse> { this.requests.push(command); this.requestTimeouts.push(timeoutMs); switch (command.type) { case "prompt": if (this.promptGate) await this.promptGate; if (this.promptError) throw this.promptError; if (this.promptResponseError) { return { type: "response", command: command.type, success: false, error: this.promptResponseError }; } return { type: "response", command: command.type, success: true }; case "prompt_and_wait": if (this.promptGate) await this.promptGate; if (this.promptError) throw this.promptError; return { type: "response", command: command.type, success: true }; case "cancel_prompt_admission": return { type: "response", command: command.type, success: true, data: { status: this.cancelPromptAdmissionStatus }, }; case "list": return { type: "response", command: command.type, success: true, data: { sessions: this.updateRestartSessions }, }; case "attach": if (this.attachFailures > 0) { this.attachFailures--; throw new Error("attach failed"); } if (command.activeSessionId === "active-restored" && this.restoredAttachGate) { await this.restoredAttachGate; this.restoredAttachCompleted++; } if (command.activeSessionId === "missing") { return { type: "response", command: command.type, success: false, error: "Unknown active session: missing", }; } return { type: "response", command: command.type, success: true, data: this.attachResultFactory?.(command) ?? createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }; case "get_queue": return { type: "response", command: command.type, success: true, data: { steering: ["steer"], followUp: ["follow"] }, }; case "get_connection_state": await this.connectionStateGate; return { type: "response", command: command.type, success: true, data: this.connectionStateFactory?.(command.activeSessionId) ?? createConnectionState(command.activeSessionId, "session-current"), }; case "get_messages": return { type: "response", command: command.type, success: true, data: { messages: [{ role: "user", content: "current prompt", timestamp: 4 }] }, }; case "get_rlm_children": await this.rlmChildrenGate; return { type: "response", command: command.type, success: true, data: { children: this.rlmChildren, eventSequence: this.rlmChildrenEventSequence }, }; case "get_resource_snapshot": return { type: "response", command: command.type, success: true, data: { contextFiles: [{ path: "/tmp/AGENTS.md" }], skills: [ { name: "demo-skill", description: "Demo skill", filePath: "/tmp/skills/demo-skill/SKILL.md", sourceInfo: { path: "/tmp/skills/demo-skill/SKILL.md", source: "local", scope: "project", origin: "top-level", baseDir: "/tmp/skills", }, }, ], prompts: [], extensions: [], themes: [], diagnostics: { skills: [], prompts: [], extensions: [], themes: [], }, }, }; case "replace_acp_mcp_servers": return { type: "response", command: command.type, success: true }; case "get_model_catalog": return { type: "response", command: command.type, success: true, data: { models: [getModel("openai", "gpt-5.1")], configuredProviders: ["openai"], }, }; case "get_available_models": return { type: "response", command: command.type, success: true, data: { models: [getModel("openai", "gpt-5.1")] }, }; case "get_session_context": return { type: "response", command: command.type, success: true, data: { context: { messages: [{ role: "user", content: "context prompt", timestamp: 3 }], thinkingLevel: "medium", model: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, }, }, }; case "get_session_tree": return { type: "response", command: command.type, success: true, data: { flatNodes: [ { entry: { type: "message", id: "user-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "user", content: "hello", timestamp: 1 }, }, }, ], leafId: "user-1", }, }; case "get_tool_definition": return { type: "response", command: command.type, success: true, data: { toolDefinition: { name: command.name, label: command.name, description: `${command.name} description`, promptSnippet: `${command.name} prompt`, promptGuidelines: [`Use ${command.name}`], parameters: { type: "object" }, renderShell: "self", }, }, }; case "clear_queue": return { type: "response", command: command.type, success: true, data: { steering: ["cleared"], followUp: [] }, }; case "abort_and_clear_queue": if (this.abortAndClearQueueUnknownCommand) { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: abort_and_clear_queue", }; } return { type: "response", command: command.type, success: true, data: { steering: ["aborted"], followUp: ["cleared"] }, }; case "acquire_session_input_pause": if (this.inputPauseAcquireGate) await this.inputPauseAcquireGate; return { type: "response", command: command.type, success: true, data: { pauseId: "pause-1" }, }; case "release_session_input_pause": return { type: "response", command: command.type, success: true }; case "heartbeats_list": return this.serverCapabilities.has("heartbeat_catalog") ? { type: "response", command: command.type, success: true, data: { heartbeats: [] } } : { type: "response", command: command.type, success: false, error: "Unknown daemon command: heartbeats_list", }; case "heartbeat_manage": return this.serverCapabilities.has("heartbeat_management") ? { type: "response", command: command.type, success: true, data: { heartbeat: { id: command.jobId } }, } : { type: "response", command: command.type, success: false, error: "Unknown daemon command: heartbeat_manage", }; case "list_saved_sessions": { const activeSessionId = "activeSessionId" in command ? command.activeSessionId : undefined; options.onProgress?.({ id: "daemon_test", type: "session_list_progress", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), loaded: 1, total: 2, }); options.onProgress?.({ id: "daemon_test", type: "session_list_item", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), session: { path: "/tmp/session-a.jsonl", id: "session-a", cwd: "/tmp", name: "Saved session", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-02T00:00:00.000Z", messageCount: 2, firstMessage: "hello", allMessagesText: "hello world", }, }); options.onProgress?.({ id: "daemon_test", type: "session_list_progress", command: "list_saved_sessions", ...(activeSessionId ? { activeSessionId } : {}), loaded: 2, total: 2, }); return { type: "response", command: command.type, success: true, data: { sessions: [ { path: "/tmp/session-a.jsonl", id: "session-a", cwd: "/tmp", name: "Saved session", created: "2026-01-01T00:00:00.000Z", modified: "2026-01-02T00:00:00.000Z", messageCount: 2, firstMessage: "hello", allMessagesText: "hello world", }, ], }, }; } case "wait_for_headless_completion": return { type: "response", command: command.type, success: true, data: { enabled: false, continuationsUsed: 0, turnsUsed: 0, tokensUsed: 0, limits: { maxContinuations: 0 }, }, }; case "roster_subscribe": return { type: "response", command: command.type, success: true, data: { roster: [] } }; case "wait_for_idle": case "set_scoped_models": case "rename_saved_session": case "extension_ui_response": case "detach": return { type: "response", command: command.type, success: true }; case "cancel_rlm_child": if (command.childId === "stale-daemon") { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: cancel_rlm_child", }; } return { type: "response", command: command.type, success: true, data: { cancelled: command.childId === "child-1" }, }; case "execute_bash": if (command.command === "stale-daemon") { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: execute_bash", }; } return { type: "response", command: command.type, success: true }; case "abort_bash": if (this.abortBashUnknownCommand) { return { type: "response", command: command.type, success: false, error: "Unknown daemon command: abort_bash", }; } return { type: "response", command: command.type, success: true }; case "start_side_question": return { type: "response", command: command.type, success: true }; case "delete_saved_session": return { type: "response", command: command.type, success: true, data: { ok: true, method: "trash" }, }; case "refine": return { type: "response", command: command.type, success: true, data: { id: "refine_daemon", summary: "Daemon refinement", rationale: "Test daemon refine timeout", expectedOutcome: "Refine request completes", appliedEdits: [], harnessStatePath: "/tmp/harness_state.json", }, }; case "cron_add": await this.cronAddGate; return { type: "response", command: command.type, success: true, data: { job: { id: `cron-${this.requests.filter((request) => request.type === "cron_add").length}`, status: "active", source: "cron", activeSessionId: command.activeSessionId, sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", cwd: "/tmp/project", prompt: command.prompt, schedule: { kind: "cron", expression: command.schedule }, createdAt: "2026-01-01T00:00:00.000Z", updatedAt: "2026-01-01T00:00:00.000Z", runCount: 0, }, }, }; case "switch_session": return { type: "response", command: command.type, success: false, error: "Stored session working directory does not exist: /tmp/missing\nSession file: /tmp/session.jsonl\nCurrent working directory: /tmp/current", errorInfo: { code: "missing_session_cwd", issue: { sessionFile: "/tmp/session.jsonl", sessionCwd: "/tmp/missing", fallbackCwd: "/tmp/current", }, }, }; case "import_jsonl": return { type: "response", command: command.type, success: false, error: "File not found: /tmp/not-found.jsonl", errorInfo: { code: "session_import_file_not_found", filePath: "/tmp/not-found.jsonl", }, }; default: throw new Error(`Unexpected command: ${command.type}`); } }
supportsServerCapability(capability: string): boolean { return this.serverCapabilities.has(capability); }
onMessage(listener: DaemonClientMessageListener): () => void { this.messageListeners.add(listener); return () => { this.messageListeners.delete(listener); }; }
onClose(listener: DaemonClientCloseListener): () => void { this.closeListeners.add(listener); return () => { this.closeListeners.delete(listener); }; }
emitMessage(message: DaemonOutbound): void { for (const listener of [...this.messageListeners]) { listener(message); } }
emitClose(error: Error): void { for (const listener of [...this.closeListeners]) { listener(error); } }
enableRequestRecovery(): void {}
async connect(): Promise<void> { if (this.connected) { throw new Error("Prime Agent daemon client is already connected"); } this.reconnectCount++; if (this.reconnectError) { throw this.reconnectError; } this.connected = true; }
async waitForHello(): Promise<DaemonHello> { return this.hello!; }
resetTransportForReconnect(): void { this.resetTransportCount++; this.connected = false; }
async reconnect(): Promise<void> { if (this.connected) { return; } this.reconnectCount++; if (this.reconnectError) { throw this.reconnectError; } this.connected = true; }
get isConnected(): boolean { return this.connected; }
getMessageListenerCount(): number { return this.messageListeners.size; }
getCloseListenerCount(): number { return this.closeListeners.size; }
close(): void { this.closeCount++; this.connected = false; if (this.emitCloseOnClose) { this.emitClose(new Error("Daemon socket closed")); } }
disconnectForReconnect(reason: "shutdown" | "update"): void { this.closeCount++; this.connected = false; this.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", reason)); }}
function asDaemonClient(client: FakeDaemonClient): DaemonTransportClient { return client as unknown as DaemonTransportClient;}
function createConnectionState(activeSessionId: string, sessionId: string): AgentConnectionState { return { activeSessionId, cwd: "/tmp/project", model: undefined, thinkingLevel: "medium", serviceTier: "default", availableThinkingLevels: ["minimal", "low", "medium", "high", "xhigh"], isStreaming: false, isCompacting: false, isBashRunning: false, retryAttempt: 0, steeringMode: "all", followUpMode: "one-at-a-time", sessionFile: `/tmp/${sessionId}.jsonl`, sessionId, sessionName: `${sessionId} name`, sessionDir: "/tmp/sessions", leafId: `${sessionId}-leaf`, autoCompactionEnabled: true, messageCount: 1, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, compactionCount: 0, goal: { active: false, status: "idle", tokensUsed: 0, timeUsedSeconds: 0, continuationsUsed: 0, }, scopedModels: [], activeToolNames: ["ipython"], contextUsage: undefined, };}
interface CreateAttachResultOptions { state?: AgentConnectionState; messages?: AgentMessage[]; streamingMessage?: AgentMessage; sessionContext?: DaemonAttachResult["snapshot"]["sessionContext"]; omitSessionContext?: boolean; sessionTree?: DaemonAttachResult["snapshot"]["sessionTree"]; parent?: DaemonAttachResult["snapshot"]["parent"]; children?: DaemonAttachResult["snapshot"]["children"]; replay?: DaemonAttachResult["replay"];}
function createAttachResult( activeSessionId: string, clientId: string | undefined, capabilities: readonly string[] | undefined, lastEventSequence: number, options: CreateAttachResultOptions = {},): DaemonAttachResult { const state = options.state ?? createConnectionState(activeSessionId, "session-current"); const messages = options.messages ?? []; const lastEventCursor = { generation: `generation-${activeSessionId}`, sequence: lastEventSequence }; const summary = { id: activeSessionId, activeSessionId, lifecycle: "live" as const, activity: "idle" as const, isSessionActive: state.isStreaming, sessionId: state.sessionId, cwd: "/tmp/project", isStreaming: state.isStreaming, isCompacting: false, attachedClients: 1, messageCount: messages.length, sessionActions: { queuedCount: 0, steering: [], followUps: [] }, ...(options.streamingMessage ? { streamingMessage: options.streamingMessage } : {}), }; // Slim shape: the daemon omits top-level state/messages for clients with the // "slim_attach" capability, which DaemonAgentConnection always advertises. return { protocol: DAEMON_PROTOCOL_INFO, activeSessionId, snapshot: { activeSessionId, summary, state, messages, ...(options.omitSessionContext ? {} : { sessionContext: options.sessionContext ?? ({ messages, thinkingLevel: state.thinkingLevel, serviceTier: state.serviceTier, model: state.model ? { provider: state.model.provider, modelId: state.model.id } : null, } satisfies NonNullable<DaemonAttachResult["snapshot"]["sessionContext"]>), }), sessionTree: options.sessionTree ?? { tree: [], leafId: state.leafId }, lastEventSequence, lastEventCursor, ...(options.parent ? { parent: options.parent } : {}), ...(options.children ? { children: options.children } : {}), }, replay: options.replay ?? { status: "complete", toSequence: lastEventSequence, toCursor: lastEventCursor, }, lastEventSequence, lastEventCursor, client: { id: clientId ?? "client-1", capabilities: (capabilities ?? ["attach_snapshot", "event_sequence"]).filter( (capability): capability is DaemonAttachResult["client"]["capabilities"][number] => capability === "attach_snapshot" || capability === "event_sequence" || capability === "extension_ui" || capability === "slim_attach" || capability === "chunked_snapshot", ), }, };}
function emitRlmChildUpdate( client: FakeDaemonClient, activeSessionId: string, sequence: number, child: AgentConnectionRlmChildAgentSnapshot,): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "rlm_child_update", child }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, });}
function emitSequencedQueueUpdate(client: FakeDaemonClient, activeSessionId: string, sequence: number): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] }, }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, });}
describe("DaemonAgentConnection", () => { it("falls back to the supervisor when the direct socket closes during initial attach", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async () => { directConnected = false; const error = new Error("direct attach socket closed"); for (const listener of [...closeListeners]) listener(error); throw error; }, close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct);
const connection = await DaemonAgentConnection.attach(routed, "active-1");
await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { activeSessionId: "active-1" }, }); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); await connection.dispose(); });
it("keeps serving the session on the direct link and reconnects a lost supervisor socket", async () => { const supervisor = new FakeDaemonClient(); const directRequests: DaemonCommand["type"][] = []; let closeSent = false; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: DaemonCommand) => { directRequests.push(command.type); if (!closeSent) { closeSent = true; supervisor.connected = false; supervisor.emitClose(new Error("supervisor closed during direct attach")); } return { type: "response" as const, command: command.type, success: true as const, data: command.type === "attach" ? createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12) : undefined, }; }, close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct);
const connection = await DaemonAgentConnection.attach(routed, "active-1"); await vi.waitFor(() => expect(supervisor.reconnectCount).toBe(1));
expect(routed.hasDirectTransport).toBe(true); // Held-direct recovery is control-plane only: no re-attach crosses either socket. expect(directRequests.filter((type) => type === "attach")).toHaveLength(1); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0); await connection.dispose(); });
it("rejects the fallback attach with the authoritative shutdown instead of parking it", async () => { const supervisor = new FakeDaemonClient(); const attachOptions: (DaemonClientRequestOptions | undefined)[] = []; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "attach") { attachOptions.push(options); throw new Error("supervisor socket is gone"); } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async () => { // The authoritative stop lands while the direct link is still up; then the direct attach dies. supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); directConnected = false; const error = new Error("direct attach socket closed"); for (const listener of [...closeListeners]) listener(error); throw error; }, close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct);
await expect(DaemonAgentConnection.attach(routed, "active-1")).rejects.toThrow("Reason: shutdown"); expect(attachOptions).toEqual([{ recoverable: false }]); });
it("absorbs a direct loss inside the held roster re-attach into a session-plane reattach", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; let rosterSubscribes = 0; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "roster_subscribe") { rosterSubscribes++; if (rosterSubscribes === 2) { directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died")); } } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await connection.subscribeAgentRoster(() => {}); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.hello = { ...supervisor.hello! };
supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost"));
await vi.waitFor(() => expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(1)); expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); await connection.dispose(); });
it("never emits connected after a shutdown that lands inside the held roster re-attach", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); let rosterSubscribes = 0; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "roster_subscribe") { rosterSubscribes++; if (rosterSubscribes === 2) { supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); } } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await connection.subscribeAgentRoster(() => {}); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.hello = { ...supervisor.hello! };
supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost"));
await vi.waitFor(() => expect(events.some((event) => event.type === "closed")).toBe(true)); expect(events.filter((event) => event.type === "closed")).toHaveLength(1); expect(events.some((event) => event.type === "connection_status" && event.status === "connected")).toBe(false); await connection.dispose(); });
it("stands down for update restoration when the update close lands inside the held roster re-attach", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; let rosterSubscribes = 0; const originalRequest = supervisor.request.bind(supervisor); supervisor.request = async (command: DaemonCommand, timeoutMs?: number, options?: DaemonClientRequestOptions) => { if (command.type === "roster_subscribe") { rosterSubscribes++; if (rosterSubscribes === 2) { supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); } } return originalRequest(command, timeoutMs ?? 30000, options ?? {}); }; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); // A short deadline would abort the restore if the held loop wrongly stayed its owner. const connection = await DaemonAgentConnection.attach(routed, "active-1", { reconnectTimeoutMs: 30 }); await connection.subscribeAgentRoster(() => {}); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.hello = { ...supervisor.hello! };
supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost"));
await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true), { timeout: 5000, }); expect(events.find((event) => event.type === "session_resynced")).toMatchObject({ snapshot: { state: { activeSessionId: "active-restored" } }, }); expect(events.some((event) => event.type === "closed")).toBe(false); await connection.dispose(); });
it("rebinds the roster subscription onto the recovered supervisor socket while the direct link holds", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("agent_roster"); const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); await connection.subscribeAgentRoster(() => {}); // A reconnect delivers a fresh hello object; the roster store keys its subscription on it. supervisor.hello = { ...supervisor.hello! };
supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket lost"));
await vi.waitFor(() => expect(supervisor.requests.filter((request) => request.type === "roster_subscribe")).toHaveLength(2), ); expect(routed.hasDirectTransport).toBe(true); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0); await connection.dispose(); });
it("resets a half-open supervisor handshake without closing the healthy direct socket", async () => { const supervisor = new FakeDaemonClient(); const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); let helloAttempts = 0; supervisor.waitForHello = vi.fn(async () => { helloAttempts++; if (helloAttempts === 1) throw new Error("hello timed out on half-open socket"); return supervisor.hello!; }); supervisor.connected = false; supervisor.emitClose(new Error("supervisor socket closed"));
await vi.waitFor(() => expect(events.some((event) => event.type === "connection_status" && event.status === "connected")).toBe(true), );
expect(supervisor.reconnectCount).toBe(2); expect(supervisor.resetTransportCount).toBe(1); expect(routed.hasDirectTransport).toBe(true); await connection.dispose(); });
it("keeps a parent direct transport intact when watching another session", async () => { const supervisor = new FakeDaemonClient(); const directRequests: DaemonCommand[] = []; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: DaemonCommand) => { directRequests.push(command); if (command.type === "attach") { return { type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }; } return { type: "response" as const, command: command.type, success: true as const }; }, close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const parent = await DaemonAgentConnection.attach(routed, "source-active");
const watcher = await parent.watchSession("target-active");
expect(watcher).toBeDefined(); await expect(parent.getState()).resolves.toMatchObject({ activeSessionId: "source-active" }); expect(routed.hasDirectTransport).toBe(true); expect( directRequests.some((request) => request.type === "attach" && request.activeSessionId === "target-active"), ).toBe(false); await watcher?.close(); await parent.dispose(); });
it("keeps a clean daemon stop terminal even while the direct link is healthy", async () => { const supervisor = new FakeDaemonClient(); const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1", { recoverDaemon: async () => {}, }); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); });
supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); await vi.waitFor(() => expect(events.some((event) => event.type === "closed")).toBe(true));
expect(events.find((event) => event.type === "closed")).toMatchObject({ error: expect.stringContaining("daemon shut down"), }); expect(supervisor.reconnectCount).toBe(0); await connection.dispose(); });
it("routes an update close through update restoration and drops the stale direct link", async () => { const supervisor = new FakeDaemonClient(); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const restored = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe(async (event) => { if (event.type === "session_resynced") resolve(event); }); });
supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update"));
await expect(restored).resolves.toMatchObject({ type: "session_resynced" }); expect(routed.hasDirectTransport).toBe(false); await connection.dispose(); });
it("falls back to the supervisor when the direct socket dies mid-session without recoverDaemon", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); supervisor.serverCapabilities.add("session_input_pause"); const pause = await connection.acquireSessionInputPause("lease-1");
directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died"));
await vi.waitFor(() => expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(1)); expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(1); // The fence died with the direct link; its holder learns on release while the session lives on. await expect(pause.release()).rejects.toThrow("invalidated by a daemon reconnect"); await connection.dispose(); });
it("keeps a shutdown during initial direct attach terminal instead of respawning the daemon", async () => { const supervisor = new FakeDaemonClient(); const recoverDaemon = vi.fn(async () => {}); let closeSent = false; const direct = { isConnected: true, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: () => () => {}, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => { if (!closeSent) { closeSent = true; supervisor.connected = false; supervisor.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); } return { type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }; }, close: () => {}, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct);
// No listener can exist during construction, so the terminal close must reject the attach. await expect(DaemonAgentConnection.attach(routed, "active-1", { recoverDaemon })).rejects.toThrow( "Reason: shutdown", ); await new Promise((resolveSettle) => setImmediate(resolveSettle)); expect(recoverDaemon).not.toHaveBeenCalled(); expect(supervisor.reconnectCount).toBe(0); });
it("takes update restoration when the direct link closes for an update while a pause is held", async () => { const supervisor = new FakeDaemonClient(); supervisor.serverCapabilities.add("session_input_pause"); supervisor.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); }); await connection.acquireSessionInputPause("lease-1");
directConnected = false; for (const listener of [...closeListeners]) { listener(new DaemonSocketClosedError("/tmp/worker.sock", "update")); }
await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true)); expect(events.some((event) => event.type === "closed")).toBe(false); await connection.dispose(); });
it("outlives the reconnect deadline while the direct link streams, then bounds recovery once it dies", async () => { const supervisor = new FakeDaemonClient(); const closeListeners = new Set<(error: Error) => void>(); let directConnected = true; const direct = { get isConnected() { return directConnected; }, hello: supervisor.hello, supportsServerCapability: (capability: string) => supervisor.supportsServerCapability(capability), onMessage: () => () => {}, onClose: (listener: (error: Error) => void) => { closeListeners.add(listener); return () => closeListeners.delete(listener); }, request: async (command: Extract<DaemonCommand, { type: "attach" }>) => ({ type: "response" as const, command: "attach" as const, success: true as const, data: createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }), close: () => { directConnected = false; }, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct); const connection = await DaemonAgentConnection.attach(routed, "active-1", { reconnectTimeoutMs: 60 }); const events: AgentConnectionEvent[] = []; connection.subscribe(async (event) => { events.push(event); });
supervisor.connected = false; supervisor.reconnectError = new Error("supervisor is down"); supervisor.emitClose(new Error("supervisor socket lost")); // Three failed control-plane attempts span well past the 60ms deadline. await vi.waitFor(() => expect(supervisor.resetTransportCount).toBeGreaterThanOrEqual(3));
expect(events.some((event) => event.type === "closed")).toBe(false); expect(supervisor.requests.filter((request) => request.type === "attach")).toHaveLength(0);
directConnected = false; for (const listener of [...closeListeners]) listener(new Error("direct worker socket died"));
await vi.waitFor(() => expect(events.some((event) => event.type === "closed")).toBe(true), { timeout: 8000 }); expect(events.find((event) => event.type === "closed")).toMatchObject({ error: expect.stringContaining("Daemon reconnection failed"), }); expect(events.filter((event) => event.type === "session_resynced")).toHaveLength(0); await connection.dispose(); }, 10_000);
it("keeps the source direct socket when a cross-worker reattach is rejected", async () => { const supervisor = new FakeDaemonClient(); supervisor.request = vi.fn(async (command: DaemonCommand) => ({ type: "response" as const, command: command.type, success: false as const, error: "target disappeared", })); const close = vi.fn(); const direct = { isConnected: true, onMessage: () => () => {}, onClose: () => () => {}, close, } as unknown as DaemonWorkerClient; const routed = new DaemonRoutedClient(asDaemonClient(supervisor), direct);
const response = await routed.request({ type: "reattach", activeSessionId: "source-active", targetActiveSessionId: "missing-target", });
expect(response.success).toBe(false); expect(close).not.toHaveBeenCalled(); expect(routed.hasDirectTransport).toBe(true); routed.close(); });
it("carries an opt-out-only telemetry policy on attach", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { telemetryDisabled: true, });
await connection.attach();
expect(fakeClient.requests[0]).toMatchObject({ type: "attach", activeSessionId: "active-1", telemetryDisabled: true, }); });
it.each([true, false])("capability-gates owned-session recovery context: %s", async (supported) => { const fakeClient = new FakeDaemonClient(); if (supported) fakeClient.serverCapabilities.add("owned_session_recovery_context"); const recoveryConfig = { cwd: "/tmp/fresh-owner" }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-owned", { ownedSession: true, ownedSessionRecoveryConfig: recoveryConfig, });
await connection.attach();
const request = fakeClient.requests[0]; if (supported) expect(request).toMatchObject({ recoveryConfig }); else expect(request).not.toHaveProperty("recoveryConfig"); });
it("forwards queueIfBusy for prompt admission", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.prompt("queued input", { streamingBehavior: "followUp", queueIfBusy: true });
expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: "queued input", streamingBehavior: "followUp", queueIfBusy: true, }); });
it("forwards hidden custom prompt metadata", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const customMessage = createManualContinuationMessage(1);
await connection.prompt(MANUAL_CONTINUATION_PROMPT, { customMessage, expandPromptTemplates: false, });
expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: MANUAL_CONTINUATION_PROMPT, customMessage, expandPromptTemplates: false, }); });
it("forwards signal-backed prompts with a unique cancellable admission id", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController();
const customMessage = createManualContinuationMessage(1); await connection.prompt("startup", { signal: abort.signal, queueIfBusy: true, customMessage, expandPromptTemplates: false, });
expect(fakeClient.requests.at(-1)).toMatchObject({ type: "prompt", activeSessionId: "active-1", message: "startup", queueIfBusy: true, customMessage, expandPromptTemplates: false, admissionId: expect.stringMatching(/^prompt-admission:/), }); });
it("cancels rejected signal-backed admission and removes its abort listener", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("transport failed"); fakeClient.cancelPromptAdmissionStatus = "cancelled"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); const add = vi.spyOn(abort.signal, "addEventListener"); const remove = vi.spyOn(abort.signal, "removeEventListener");
await expect(connection.prompt("startup", { signal: abort.signal })).rejects.toMatchObject({ message: "transport failed", cancelled: true, }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt", "cancel_prompt_admission"]); expect(add).toHaveBeenCalledOnce(); expect(remove).toHaveBeenCalledOnce(); expect(remove.mock.calls[0]?.[1]).toBe(add.mock.calls[0]?.[1]); });
it("preserves a definitive prompt rejection when the signal remains live", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptResponseError = "session rejected prompt"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toEqual( new Error("session rejected prompt"), ); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt"]); });
it("sends zero requests for a pre-aborted prompt", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController(); abort.abort();
await expect(connection.prompt("startup", { signal: abort.signal })).rejects.toMatchObject({ status: "cancelled", }); expect(fakeClient.requests).toEqual([]); });
it("accepts a successful prompt response when cancellation reports unknown", async () => { const fakeClient = new FakeDaemonClient(); let releasePrompt = () => {}; fakeClient.promptGate = new Promise<void>((resolve) => { releasePrompt = resolve; }); fakeClient.cancelPromptAdmissionStatus = "unknown"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController();
const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); expect(fakeClient.requests.find((request) => request.type === "cancel_prompt_admission")).not.toHaveProperty( "cancelOwned", ); releasePrompt();
await expect(prompt).resolves.toBeUndefined(); });
it("requests owned prompt cancellation when the daemon advertises it", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("owned_prompt_cancellation"); let releasePrompt = () => {}; fakeClient.promptGate = new Promise<void>((resolve) => { releasePrompt = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController();
const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); expect(fakeClient.requests.find((request) => request.type === "cancel_prompt_admission")).toMatchObject({ cancelOwned: true, }); releasePrompt(); await expect(prompt).resolves.toBeUndefined(); });
it("preserves a definitive prompt rejection when cancellation reports owned", async () => { const fakeClient = new FakeDaemonClient(); let releasePrompt = () => {}; fakeClient.promptGate = new Promise<void>((resolve) => { releasePrompt = resolve; }); fakeClient.promptResponseError = "post-ownership validation failed"; fakeClient.cancelPromptAdmissionStatus = "owned"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const abort = new AbortController();
const prompt = connection.prompt("startup", { signal: abort.signal }); abort.abort(); await vi.waitFor(() => expect(fakeClient.requests.map((request) => request.type)).toContain("cancel_prompt_admission"), ); releasePrompt();
await expect(prompt).rejects.toThrow("post-ownership validation failed"); });
it("treats a lost prompt response plus unknown cancellation as uncertain", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("lost response"); fakeClient.cancelPromptAdmissionStatus = "unknown"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toMatchObject({ message: "lost response", status: "unknown", cancelled: false, }); });
it("translates unsupported admission cancellation into an admission error", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new DaemonCapabilityUnavailableError("prompt", "prompt_admission_cancellation"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await expect(connection.prompt("startup", { signal: new AbortController().signal })).rejects.toMatchObject({ status: "unsupported", }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["prompt"]); });
it("uses cancellable admission for promptAndWait", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.promptError = new Error("lost response"); fakeClient.cancelPromptAdmissionStatus = "cancelled"; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await expect(connection.promptAndWait("wait", { signal: new AbortController().signal })).rejects.toMatchObject({ status: "cancelled", }); expect(fakeClient.requests.map((request) => request.type)).toEqual([ "prompt_and_wait", "cancel_prompt_admission", ]); });
it("uses fleet heartbeat scope for residents and session scope for owned workers", async () => { const residentClient = new FakeDaemonClient(); residentClient.serverCapabilities.add("heartbeat_catalog"); const resident = new DaemonAgentConnection(asDaemonClient(residentClient), "resident-1"); await resident.listHeartbeats();
const ownedClient = new FakeDaemonClient(); ownedClient.serverCapabilities.add("heartbeat_catalog"); const owned = new DaemonAgentConnection(asDaemonClient(ownedClient), "owned-1", { ownedSession: true }); await owned.listHeartbeats();
expect(residentClient.requests.at(-1)).toEqual(expect.objectContaining({ type: "heartbeats_list" })); expect(residentClient.requests.at(-1)).not.toHaveProperty("activeSessionId"); expect(ownedClient.requests.at(-1)).toEqual( expect.objectContaining({ type: "heartbeats_list", activeSessionId: "owned-1" }), ); });
it("serializes concurrent owned-session promotion commands", async () => { const fakeClient = new FakeDaemonClient(); let releaseFirst = () => {}; fakeClient.cronAddGate = new Promise<void>((resolve) => { releaseFirst = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { ownedSession: true, });
const first = connection.addCronJob("0 * * * *", "first"); const second = connection.addCronJob("30 * * * *", "second"); await vi.waitFor(() => { expect(fakeClient.requests.filter((request) => request.type === "cron_add")).toHaveLength(1); }); releaseFirst(); await Promise.all([first, second]);
const commands = fakeClient.requests.filter( (command): command is Extract<DaemonCommand, { type: "cron_add" }> => command.type === "cron_add", ); expect(commands.map((command) => command.promoteOwnedSession)).toEqual([true, false]); });
it("gates side-question follow-up transcripts on the daemon capability", async () => { const oldDaemonClient = new FakeDaemonClient(); const oldConnection = new DaemonAgentConnection(asDaemonClient(oldDaemonClient), "active-original");
// First questions carry no transcript and must keep working on old daemons. await oldConnection.startSideQuestion("turn-1", "What changed?"); expect(oldDaemonClient.requests.map((request) => request.type)).toEqual(["start_side_question"]);
// Follow-ups would silently lose their side context on an old daemon. await expect( oldConnection.startSideQuestion("turn-2", "And then?", [{ question: "What changed?", answer: "The parser." }]), ).rejects.toThrow("older build without side-conversation follow-ups"); expect(oldDaemonClient.requests).toHaveLength(1);
const newDaemonClient = new FakeDaemonClient(); newDaemonClient.serverCapabilities.add("side_question_transcript"); const newConnection = new DaemonAgentConnection(asDaemonClient(newDaemonClient), "active-original");
await newConnection.startSideQuestion("turn-2", "And then?", [ { question: "What changed?", answer: "The parser." }, ]); const sent = newDaemonClient.requests.find( (command): command is Extract<DaemonCommand, { type: "start_side_question" }> => command.type === "start_side_question", ); expect(sent?.previousTurns).toEqual([{ question: "What changed?", answer: "The parser." }]); });
it("gates transient bash on the daemon capability", async () => { const oldDaemonClient = new FakeDaemonClient(); const oldConnection = new DaemonAgentConnection(asDaemonClient(oldDaemonClient), "active-original");
// Regular bash keeps working on old daemons. await oldConnection.executeBash("ls"); expect(oldDaemonClient.requests.map((request) => request.type)).toEqual(["execute_bash"]);
// A transient run on an old daemon would be recorded into the session. await expect(oldConnection.executeBash("ls", { transient: true })).rejects.toThrow( "older build without side-conversation bash", ); expect(oldDaemonClient.requests).toHaveLength(1);
const newDaemonClient = new FakeDaemonClient(); newDaemonClient.serverCapabilities.add("transient_bash"); const newConnection = new DaemonAgentConnection(asDaemonClient(newDaemonClient), "active-original");
await newConnection.executeBash("ls", { excludeFromContext: true, transient: true, runId: "side-run-1" }); const sent = newDaemonClient.requests.find( (command): command is Extract<DaemonCommand, { type: "execute_bash" }> => command.type === "execute_bash", ); expect(sent).toMatchObject({ command: "ls", excludeFromContext: true, transient: true, runId: "side-run-1" }); });
it("degrades an unavailable heartbeat catalog without sending an unsupported command", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original");
await expect(connection.listHeartbeats()).resolves.toEqual([]); expect(fakeClient.requests).toEqual([]); await expect(connection.manageHeartbeat("active-original", "job-1", "pause")).rejects.toThrow( "requires a newer Prime Agent daemon", ); expect(fakeClient.requests).toEqual([]); });
it("reattaches an open window to its restored session after an update restart", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.emitCloseOnClose = true; const restoredMessages: AgentMessage[] = [{ role: "user", content: "restored prompt", timestamp: 2 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, "session-current"), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const events: AgentConnectionEvent[] = []; const restored = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { events.push(event); if (event.type === "session_resynced") { resolve(event); } }); }); await connection.attach();
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); await vi.waitFor(() => { expect(fakeClient.closeCount).toBe(1); expect(fakeClient.reconnectCount).toBe(1); });
await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-current" }, messages: restoredMessages, }, }); expect(fakeClient.reconnectCount).toBe(1); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", resumeCursor: undefined, }); await vi.waitFor(() => { expect(events).toEqual([ expect.objectContaining({ type: "connection_status", status: "reconnecting" }), expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ activeSessionId: "active-restored" }), }), }), { type: "connection_status", status: "connected" }, ]); }); });
it("reattaches when an update socket close arrives before the session notice", async () => { const fakeClient = new FakeDaemonClient(); const restoredMessages: AgentMessage[] = [{ role: "user", content: "restored prompt", timestamp: 2 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, "session-current"), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const restored = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update"));
await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-current" }, messages: restoredMessages, }, }); expect(fakeClient.reconnectCount).toBe(1); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); });
it("coordinates one transport reconnect across connections sharing a daemon client", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.emitCloseOnClose = true; fakeClient.updateRestartSessions = [ { id: "restored-a", activeSessionId: "restored-a", sessionId: "session-a", sessionFile: "/tmp/session-a.jsonl", }, { id: "restored-b", activeSessionId: "restored-b", sessionId: "session-b", sessionFile: "/tmp/session-b.jsonl", }, ]; const sessionIds: Record<string, string> = { "active-a": "session-a", "active-b": "session-b", "restored-a": "session-a", "restored-b": "session-b", }; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, sessionIds[command.activeSessionId]!), }); const connectionA = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-a"); const connectionB = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-b"); await connectionA.attach(); await connectionB.attach(); const restoredA = new Promise<AgentConnectionEvent>((resolve) => { connectionA.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); const restoredB = new Promise<AgentConnectionEvent>((resolve) => { connectionB.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); });
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-a", reason: "update" });
await expect(Promise.all([restoredA, restoredB])).resolves.toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-a" }), }), }), expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-b" }), }), }), ]); expect(fakeClient.closeCount).toBe(1); expect(fakeClient.reconnectCount).toBe(1); });
it("does not reconnect after an explicit shutdown session close", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "shutdown", }); fakeClient.emitClose(new Error("Daemon socket closed")); await Promise.resolve();
expect(fakeClient.reconnectCount).toBe(0); expect(closedEvents).toHaveLength(1); expect(closedEvents[0]).toMatchObject({ type: "closed", error: expect.stringContaining("The Prime Agent daemon shut down while this window was attached."), }); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("Session ID: session-current."); expect(closedError).toContain("Session file: /tmp/session-current.jsonl."); expect(closedError).toContain("Diagnostic log:"); });
it("does not infer an update when shutdown closes the socket before the session notice", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown")); await Promise.resolve();
expect(fakeClient.reconnectCount).toBe(0); expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("The Prime Agent daemon shut down while this window was attached."); });
it.each([ ["killed", "The daemon stopped this agent session."], ["completed", "The daemon closed this agent session after it completed."], ["replaced", "The daemon replaced this agent session with another session."], ] as const)("explains a %s session close instead of exposing the raw reason", async (reason, explanation) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason }); await Promise.resolve();
expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain(explanation); expect(closedError).not.toBe(reason); expect(closedError).toContain("Session ID: session-current."); });
it("adds recovery and diagnostic context to unexpected daemon disconnects", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitClose(new Error("ECONNRESET")); await Promise.resolve();
expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("Lost connection to the Prime Agent daemon. Cause: ECONNRESET"); expect(closedError).toContain("restart Prime Agent or reopen the session from Agents View"); expect(closedError).toContain("Session file: /tmp/session-current.jsonl."); expect(closedError).toContain("Diagnostic log:"); });
it("reports an unannounced clean socket close without treating it as an update", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock")); await Promise.resolve();
expect(fakeClient.reconnectCount).toBe(0); expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain("Lost connection to the Prime Agent daemon."); });
it("does not emit a restored session after disposal begins", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; let releaseRestoredAttach: (() => void) | undefined; fakeClient.restoredAttachGate = new Promise<void>((resolve) => { releaseRestoredAttach = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.waitFor(() => { expect( fakeClient.requests.some( (request) => request.type === "attach" && request.activeSessionId === "active-restored", ), ).toBe(true); }); await connection.dispose(); releaseRestoredAttach?.(); await vi.waitFor(() => { expect(fakeClient.restoredAttachCompleted).toBe(1); }); for (let flush = 0; flush < 5; flush++) { await Promise.resolve(); }
expect(events).toEqual([expect.objectContaining({ type: "connection_status", status: "reconnecting" })]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "detach", activeSessionId: "active-restored" }); });
it("returns to normal close handling after update restoration times out", async () => { vi.useFakeTimers(); try { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("daemon unavailable"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const closedEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { if (event.type === "closed") { closedEvents.push(event); } }); await connection.attach();
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); await vi.advanceTimersByTimeAsync(120100);
expect(closedEvents).toHaveLength(1); const closedError = closedEvents[0]?.type === "closed" ? closedEvents[0].error : undefined; expect(closedError).toContain( "The Prime Agent daemon restarted for an update, but this window could not reconnect", ); expect(closedError).toContain("Last error: daemon unavailable"); expect(closedError).toContain("restart Prime Agent and reopen it from Agents View"); expect(closedError).toContain("Session ID: session-current."); expect(closedError).toContain("Session file: /tmp/session-current.jsonl."); expect(closedError).toContain("Diagnostic log:"); const reconnectCountAfterFailure = fakeClient.reconnectCount; fakeClient.emitClose(new Error("Daemon socket closed")); await Promise.resolve();
expect(fakeClient.reconnectCount).toBe(reconnectCountAfterFailure); expect(closedEvents).toHaveLength(1); await connection.dispose(); } finally { vi.useRealTimers(); } });
it("reattaches using the replacement session identity after a session switch", async () => { const fakeClient = new FakeDaemonClient(); const restoredMessages: AgentMessage[] = [{ role: "user", content: "switched prompt", timestamp: 3 }]; fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-next", sessionFile: "/tmp/session-next.jsonl", }, ]; fakeClient.attachResultFactory = (command) => { const sessionId = command.activeSessionId === "active-restored" ? "session-next" : "session-current"; return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: createConnectionState(command.activeSessionId, sessionId), messages: command.activeSessionId === "active-restored" ? restoredMessages : [], }); }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); await connection.attach(); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-original", state: createConnectionState("active-original", "session-next"), messages: [{ role: "user", content: "switched prompt", timestamp: 2 }], });
const restored = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") { resolve(event); } }); }); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); fakeClient.emitClose(new Error("Daemon socket closed"));
await expect(restored).resolves.toMatchObject({ type: "session_resynced", snapshot: { state: { activeSessionId: "active-restored", sessionId: "session-next" }, messages: restoredMessages, }, }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "list", "attach"]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", }); });
it("loads connection state and forwards replacement snapshots through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach();
await expect(connection.getState()).resolves.toMatchObject({ activeSessionId: "active-1", cwd: "/tmp/project", sessionId: "session-current", sessionName: "session-current name", activeToolNames: ["ipython"], });
fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "session-next"), messages: [{ role: "user", content: "replacement prompt", timestamp: 1 }], }); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "other", state: createConnectionState("other", "ignored"), messages: [{ role: "user", content: "ignored", timestamp: 2 }], });
expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); expect(events).toEqual([ { type: "session_replaced", state: expect.objectContaining({ activeSessionId: "active-1", sessionId: "session-next", sessionName: "session-next name", activeToolNames: ["ipython"], }), messages: [{ role: "user", content: "replacement prompt", timestamp: 1 }], }, ]); });
it("forwards catch-up snapshots as non-destructive resync events", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); const messages: AgentMessage[] = [{ role: "user", content: "caught up", timestamp: 2 }]; const streamingMessage = { role: "assistant", content: [{ type: "thinking", thinking: "Still reasoning" }], } as AgentMessage; const snapshot = createAttachResult("active-1", "client-1", undefined, 13, { state: { ...createConnectionState("active-1", "session-current"), isStreaming: true }, messages, streamingMessage, }).snapshot;
fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, });
expect(events).toEqual([ { type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-current" }), messages, streamingMessage, lastEventSequence: 13, }), }, ]); });
it("exposes attach snapshots as the initial connection snapshot", async () => { const fakeClient = new FakeDaemonClient(); const snapshotMessage: AgentMessage = { role: "user", content: "snapshot prompt", timestamp: 1 }; const messages: AgentMessage[] = [snapshotMessage]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 15, { state: createConnectionState(command.activeSessionId, "session-snapshot"), messages, sessionContext: { messages, thinkingLevel: "medium", serviceTier: "default", model: null, }, sessionTree: { tree: [ { entry: { type: "message", id: "user-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", message: snapshotMessage, }, children: [], }, ], leafId: "user-1", }, parent: { activeSessionId: "parent-active", sessionId: "parent-session", nodeId: "parent-node", childId: "child-1", }, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.attach();
await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { activeSessionId: "active-1", sessionId: "session-snapshot", }, messages, sessionContext: { messages, thinkingLevel: "medium", model: null, }, sessionTree: { leafId: "user-1", }, parent: { activeSessionId: "parent-active", sessionId: "parent-session", nodeId: "parent-node", childId: "child-1", }, lastEventSequence: 15, replay: { status: "complete", toSequence: 15, }, }); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-snapshot", }); await expect(connection.getMessages()).resolves.toEqual(messages); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]); });
it("preserves the in-flight assistant message when refreshing a stale snapshot", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const streamingMessage: AgentMessage = { role: "assistant", content: [{ type: "text", text: "partial response" }], api: "test-api", provider: "test-provider", model: "test-model", usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, }, stopReason: "stop", timestamp: 2, }; fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "message_start", message: streamingMessage }, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, });
await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ streamingMessage }); expect(fakeClient.requests.map((request) => request.type)).toEqual([ "attach", "get_connection_state", "get_messages", "get_session_context", ]); });
it("does not stamp independently fetched snapshot data with an event cursor observed mid-fetch", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitSequencedQueueUpdate(fakeClient, "active-1", 13); let releaseState!: () => void; fakeClient.connectionStateGate = new Promise<void>((resolveState) => { releaseState = resolveState; });
const snapshotPromise = connection.getInitialSnapshot(); await Promise.resolve(); emitSequencedQueueUpdate(fakeClient, "active-1", 14); releaseState(); const snapshot = await snapshotPromise;
expect(snapshot.lastEventCursor).toEqual({ generation: "generation-active-1", sequence: 13 }); await connection.getState(); expect(fakeClient.requests.filter((request) => request.type === "get_connection_state")).toHaveLength(2); });
it("times out an attach whose streamed snapshot never completes", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => { const result = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); return { ...result, snapshotStream: { id: "snapshot-stalled", messageCount: 0, targetChunkBytes: 512 * 1024 }, }; }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { snapshotTimeoutMs: 10, });
await expect(connection.attach()).rejects.toThrow("Timed out waiting for snapshot snapshot-stalled"); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "snapshot-stalled", snapshot: createAttachResult("active-1", "client-1", undefined, 12).snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "snapshot-stalled", chunkCount: 0, lastEventSequence: 12, }); expect((connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size).toBe(0); });
it("rejects one failed snapshot without interrupting another session on the shared client", async () => { const fakeClient = new FakeDaemonClient(); const sibling = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-2"); await sibling.attach(); const siblingEvents: AgentConnectionEvent[] = []; sibling.subscribe((event) => { siblingEvents.push(event); }); fakeClient.attachResultFactory = (command) => { const result = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); if (command.activeSessionId !== "active-1") { return result; } const { messages: _messages, ...snapshot } = result.snapshot; queueMicrotask(() => { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: command.activeSessionId, snapshotId: "snapshot-failed", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: command.activeSessionId, snapshotId: "snapshot-failed", error: "snapshot encoder failed", }); }); return { ...result, snapshot: { ...result.snapshot, messages: [] }, snapshotStream: { id: "snapshot-failed", messageCount: 0, targetChunkBytes: 512 * 1024 }, }; }; const failed = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await expect(failed.attach()).rejects.toThrow("snapshot encoder failed"); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1));
expect(siblingEvents[0]).toMatchObject({ type: "session_event", event: { type: "session_action_update" } }); expect(fakeClient.closeCount).toBe(0); await failed.dispose(); await sibling.dispose(); });
it("assembles chunked attach snapshots even when chunks arrive before the attach response continuation", async () => { const fakeClient = new FakeDaemonClient(); const messages: AgentMessage[] = [ { role: "user", content: "first", timestamp: 1 }, { role: "user", content: "second", timestamp: 2 }, ]; fakeClient.attachResultFactory = (command) => { const full = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 23, { state: createConnectionState(command.activeSessionId, "session-streamed"), messages, }); const { messages: _messages, ...snapshotHeader } = full.snapshot; queueMicrotask(() => { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", snapshot: snapshotHeader, messageCount: messages.length, targetChunkBytes: 512 * 1024, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", index: 0, messages: [messages[0]!], }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", index: 1, messages: [messages[1]!], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: command.activeSessionId, snapshotId: "snapshot-streamed", chunkCount: 2, lastEventSequence: 23, }); }); return { ...full, snapshot: { ...full.snapshot, messages: [] }, snapshotStream: { id: "snapshot-streamed", messageCount: messages.length, targetChunkBytes: 512 * 1024, }, }; }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); });
await connection.attach();
await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { sessionId: "session-streamed" }, messages, lastEventSequence: 23, }); expect(events).toEqual([]); });
it("distinguishes chunked catch-up snapshots from runtime replacements", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); });
const emitSnapshot = ( purpose: "replacement" | "resync", snapshotId: string, sequence: number, sessionId: string, ) => { const messages: AgentMessage[] = [{ role: "user", content: purpose, timestamp: sequence }]; const full = createAttachResult("active-1", "client-1", undefined, sequence, { state: createConnectionState("active-1", sessionId), messages, }); const { messages: _messages, ...snapshot } = full.snapshot; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId, snapshot, messageCount: messages.length, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId, index: 0, messages, }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId, chunkCount: 1, lastEventSequence: sequence, lastEventCursor: { generation: "generation-active-1", sequence }, }); };
emitSnapshot("resync", "snapshot-resync", 13, "session-current"); emitSnapshot("replacement", "snapshot-replacement", 14, "session-next"); await vi.waitFor(() => expect(events).toHaveLength(2));
expect(events).toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ state: expect.objectContaining({ sessionId: "session-current" }), }), }), expect.objectContaining({ type: "session_replaced", state: expect.objectContaining({ sessionId: "session-next" }), }), ]); expect((connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size).toBe(0); });
it.each(["replacement", "resync"] as const)( "recovers a failed chunked %s snapshot without interrupting a sibling session", async (purpose) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const sibling = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-2"); await Promise.all([connection.attach(), sibling.attach()]); const events: AgentConnectionEvent[] = []; const siblingEvents: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); sibling.subscribe((event) => { siblingEvents.push(event); }); const recoveredSessionId = purpose === "replacement" ? "session-next" : "session-current"; fakeClient.connectionStateFactory = (activeSessionId) => createConnectionState( activeSessionId, activeSessionId === "active-1" ? recoveredSessionId : "session-sibling", ); fakeClient.requests.length = 0; const snapshotId = `snapshot-failed-${purpose}`; const full = createAttachResult("active-1", "client-1", undefined, 13, { state: createConnectionState("active-1", recoveredSessionId), }); const { messages: _messages, ...snapshot } = full.snapshot; if (purpose === "replacement") { fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", recoveredSessionId), messages: [], snapshotFollows: true, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); } fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId, snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: "active-1", snapshotId, error: `${purpose} snapshot failed`, });
await vi.waitFor(() => expect(events).toHaveLength(1)); if (purpose === "replacement") { expect(events[0]).toMatchObject({ type: "session_replaced", state: { sessionId: recoveredSessionId }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], }); } else { expect(events[0]).toMatchObject({ type: "session_resynced", snapshot: { state: { sessionId: recoveredSessionId }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], }, }); } expect(fakeClient.requests.map((request) => request.type)).toEqual([ "get_connection_state", "get_messages", "get_session_context", ]); emitSequencedQueueUpdate(fakeClient, "active-2", 13); await vi.waitFor(() => expect(siblingEvents).toHaveLength(1)); expect(siblingEvents[0]).toMatchObject({ type: "session_event", event: { type: "session_action_update" } }); expect(fakeClient.closeCount).toBe(0); expect((connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size).toBe( 0, ); await connection.dispose(); await sibling.dispose(); }, );
it("keeps attach snapshots usable when the daemon omits duplicate session context", async () => { const fakeClient = new FakeDaemonClient(); const messages: AgentMessage[] = [{ role: "user", content: "snapshot prompt", timestamp: 1 }]; fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 15, { state: createConnectionState(command.activeSessionId, "session-snapshot"), messages, omitSessionContext: true, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.attach();
const snapshot = await connection.getInitialSnapshot(); expect(snapshot).toMatchObject({ state: { activeSessionId: "active-1", sessionId: "session-snapshot", }, messages, lastEventSequence: 15, }); expect(snapshot.sessionContext).toBeUndefined(); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach"]);
await expect(connection.getSessionContext()).resolves.toMatchObject({ messages: [{ role: "user", content: "context prompt", timestamp: 3 }], }); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "get_session_context"]); });
it("keeps reconnect usable from attach snapshots when replay is unavailable", async () => { const fakeClient = new FakeDaemonClient(); const reconnectedMessages: AgentMessage[] = [{ role: "user", content: "reconnected prompt", timestamp: 2 }]; let attachCount = 0; fakeClient.attachResultFactory = (command) => { attachCount++; if (attachCount === 1) { return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: createConnectionState(command.activeSessionId, "session-initial"), }); } return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 20, { state: createConnectionState(command.activeSessionId, "session-reconnected"), messages: reconnectedMessages, sessionContext: { messages: reconnectedMessages, thinkingLevel: "medium", serviceTier: "default", model: null, }, replay: { status: "unavailable", fromSequence: 14, toSequence: 20, reason: "event_replay_not_available", }, }); }; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.attach(); emitSequencedQueueUpdate(fakeClient, "active-1", 14); await connection.attach();
expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", sequence: 14, }, }); await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ state: { sessionId: "session-reconnected", }, messages: reconnectedMessages, replay: { status: "unavailable", fromSequence: 14, toSequence: 20, reason: "event_replay_not_available", }, lastEventSequence: 20, }); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-reconnected", }); await expect(connection.getMessages()).resolves.toEqual(reconnectedMessages); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "attach"]); });
it("resets a connected transport when reattach fails during supervisor recovery", async () => { const fakeClient = new FakeDaemonClient(); const recoverDaemon = vi.fn(async () => undefined); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon, reconnectTimeoutMs: 2000, }); const statuses: string[] = []; let resyncs = 0; connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } if (event.type === "session_resynced") { resyncs++; } }); await connection.attach();
fakeClient.attachFailures = 1; fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed"));
await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(recoverDaemon).toHaveBeenCalledTimes(2); expect(fakeClient.reconnectCount).toBe(2); expect(fakeClient.resetTransportCount).toBe(1); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(3); expect(resyncs).toBe(1); });
it("does not reconnect after disposal while daemon recovery is pending", async () => { const fakeClient = new FakeDaemonClient(); let finishRecovery: (() => void) | undefined; const recoverDaemon = vi.fn( () => new Promise<void>((resolve) => { finishRecovery = resolve; }), ); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon, reconnectTimeoutMs: 2000, }); await connection.attach();
fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed")); await vi.waitFor(() => expect(recoverDaemon).toHaveBeenCalledOnce()); await connection.dispose(); finishRecovery?.(); for (let flush = 0; flush < 5; flush++) { await Promise.resolve(); }
expect(fakeClient.reconnectCount).toBe(0); });
it("isolates subscriber failures during transport recovery", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 2000, }); const statuses: string[] = []; connection.subscribe(async () => { throw new Error("broken subscriber"); }); connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } }); await connection.attach();
fakeClient.connected = false; fakeClient.emitClose(new Error("Daemon socket closed"));
await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(fakeClient.reconnectCount).toBe(1); });
it("does not let a stalled subscriber block update recovery", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [ { id: "active-restored", activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ]; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original"); const statuses: string[] = []; connection.subscribe(() => new Promise<void>(() => undefined)); connection.subscribe((event) => { if (event.type === "connection_status") { statuses.push(event.status); } }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update"));
await vi.waitFor(() => expect(statuses).toEqual(["reconnecting", "connected"])); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-restored", }); });
it.each([true, false])("capability-gates authoritative child-roster reads: %s", async (supported) => { const fakeClient = new FakeDaemonClient(); if (supported) fakeClient.serverCapabilities.add("authoritative_child_roster"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
const children = connection.getRlmChildSnapshots(); if (supported) await expect(children).resolves.toEqual([]); else await expect(children).rejects.toThrow("authoritative_child_roster"); });
it("preserves a newer live child update over an in-flight roster read", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("authoritative_child_roster"); let releaseRoster!: () => void; fakeClient.rlmChildrenGate = new Promise<void>((resolve) => { releaseRoster = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
const children = connection.getRlmChildSnapshots(); emitRlmChildUpdate(fakeClient, "active-1", 13, { id: "child-live", label: "live child", status: "running", sessionDir: "/tmp/child-live", }); releaseRoster();
await expect(children).resolves.toEqual([expect.objectContaining({ id: "child-live", status: "running" })]); });
it("keeps the live child roster after an empty attach snapshot", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { children: [] }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.attach(); emitRlmChildUpdate(fakeClient, "active-1", 13, { id: "child-live", label: "live child", status: "running", sessionDir: "/tmp/child-live", });
await expect(connection.getInitialSnapshot()).resolves.toMatchObject({ children: [expect.objectContaining({ id: "child-live", status: "running" })], }); });
it("refreshes initial snapshots after live events make the cached snapshot stale", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: createConnectionState(command.activeSessionId, "session-attached"), messages: [{ role: "user", content: "attached prompt", timestamp: 1 }], }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1");
await connection.attach(); emitSequencedQueueUpdate(fakeClient, "active-1", 13);
const snapshot = await connection.getInitialSnapshot(); expect(snapshot).toMatchObject({ state: { sessionId: "session-current", }, messages: [{ role: "user", content: "current prompt", timestamp: 4 }], sessionContext: { messages: [{ role: "user", content: "context prompt", timestamp: 3 }], }, }); // The session tree is fetched lazily (only when the tree/branch selector is // opened), so refreshing the initial snapshot must not request it. expect(snapshot.sessionTree).toBeUndefined(); expect(fakeClient.requests.map((request) => request.type)).toEqual([ "attach", "get_connection_state", "get_messages", "get_session_context", ]); });
it("ignores older sequenced events after an attach snapshot", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); });
await connection.attach(); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "session-old"), messages: [{ role: "user", content: "old prompt", timestamp: 1 }], meta: { id: "active-1:10", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 10, emittedAt: "2026-01-01T00:00:00.000Z", }, }); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "session-new"), messages: [{ role: "user", content: "new prompt", timestamp: 2 }], meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, emittedAt: "2026-01-01T00:00:00.000Z", }, });
expect(events).toEqual([ { type: "session_replaced", state: expect.objectContaining({ sessionId: "session-new", }), messages: [{ role: "user", content: "new prompt", timestamp: 2 }], }, ]); await expect(connection.getState()).resolves.toMatchObject({ sessionId: "session-new", }); });
it("sends queue commands through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_input_pause"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect(connection.getQueue()).resolves.toEqual({ steering: ["steer"], followUp: ["follow"] }); await expect(connection.clearQueue()).resolves.toEqual({ steering: ["cleared"], followUp: [] }); await expect(connection.abortAndClearQueue()).resolves.toEqual({ steering: ["aborted"], followUp: ["cleared"] }); const [inputPause, inputPauseAlias] = await Promise.all([ connection.acquireSessionInputPause("lease-1"), connection.acquireSessionInputPause("lease-1"), ]); expect(inputPauseAlias).toBe(inputPause); (connection as unknown as { activeSessionId: string }).activeSessionId = "active-2"; await inputPause.release(); (connection as unknown as { activeSessionId: string }).activeSessionId = "active-1"; await connection.waitForIdle(); const model = getModel("anthropic", "claude-sonnet-4-5"); if (!model) { throw new Error("Test model not found"); } await connection.setScopedModels([{ model, thinkingLevel: "high" }]);
expect(fakeClient.requests.map((request) => request.type)).toEqual([ "attach", "get_queue", "clear_queue", "abort_and_clear_queue", "acquire_session_input_pause", "release_session_input_pause", "wait_for_idle", "set_scoped_models", ]); expect(fakeClient.requests[1]).toMatchObject({ type: "get_queue", activeSessionId: "active-1" }); expect(fakeClient.requests[2]).toMatchObject({ type: "clear_queue", activeSessionId: "active-1" }); expect(fakeClient.requests[3]).toMatchObject({ type: "abort_and_clear_queue", activeSessionId: "active-1" }); expect(fakeClient.requests[4]).toMatchObject({ type: "acquire_session_input_pause", activeSessionId: "active-1", leaseKey: "lease-1", }); expect(fakeClient.requests[5]).toMatchObject({ type: "release_session_input_pause", activeSessionId: "active-1", pauseId: "pause-1", }); expect(fakeClient.requests[6]).toMatchObject({ type: "wait_for_idle", activeSessionId: "active-1" }); expect(fakeClient.requests[7]).toMatchObject({ type: "set_scoped_models", activeSessionId: "active-1", scopedModels: [{ model, thinkingLevel: "high" }], });
fakeClient.abortAndClearQueueUnknownCommand = true; await expect(connection.abortAndClearQueue()).rejects.toThrow( "the daemon is running an older build; restart the daemon and try again", ); });
it("fails closed when a daemon disconnect invalidates an input pause", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_input_pause"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); });
const pause = await connection.acquireSessionInputPause("lease-1"); fakeClient.disconnectForReconnect("shutdown");
await expect(pause.release()).rejects.toThrow("invalidated by a daemon reconnect"); await expect(connection.acquireSessionInputPause("lease-1")).rejects.toThrow("connection is closed"); expect(events).toContainEqual( expect.objectContaining({ type: "closed", error: expect.stringContaining("fence was invalidated"), }), ); });
it("releases an input pause whose acquisition resolves after disconnect", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_input_pause"); let releaseAcquire!: () => void; fakeClient.inputPauseAcquireGate = new Promise<void>((resolve) => { releaseAcquire = resolve; }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
const acquisition = connection.acquireSessionInputPause("lease-1"); await vi.waitFor(() => expect(fakeClient.requests.some((request) => request.type === "acquire_session_input_pause")).toBe(true), ); fakeClient.disconnectForReconnect("shutdown"); releaseAcquire();
await expect(acquisition).rejects.toThrow("acquisition was invalidated by a daemon reconnect"); expect(fakeClient.requests.filter((request) => request.type === "release_session_input_pause")).toHaveLength(1); });
it("capability-gates the strong RLM completion barrier", async () => { const oldDaemonClient = new FakeDaemonClient(); const oldConnection = new DaemonAgentConnection(asDaemonClient(oldDaemonClient), "active-old"); await oldConnection.attach(); await oldConnection.waitForHeadlessCompletion(); expect(oldDaemonClient.requests.at(-1)).toEqual({ type: "wait_for_headless_completion", activeSessionId: "active-old", }); await expect(oldConnection.waitForHeadlessCompletion({ waitForRlmQuiescence: true })).rejects.toThrow( "without RLM quiescence barriers", ); expect(oldDaemonClient.requests.map((request) => request.type)).toEqual([ "attach", "wait_for_headless_completion", ]);
const newDaemonClient = new FakeDaemonClient(); newDaemonClient.serverCapabilities.add("rlm_quiescence_barrier"); const newConnection = new DaemonAgentConnection(asDaemonClient(newDaemonClient), "active-new"); await newConnection.attach(); await newConnection.waitForHeadlessCompletion({ waitForRlmQuiescence: true }); expect(newDaemonClient.requests.at(-1)).toMatchObject({ type: "wait_for_headless_completion", activeSessionId: "active-new", waitForRlmQuiescence: true, }); });
it("cancels rlm child runs through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect(connection.cancelRlmChild("child-1")).resolves.toBe(true); await expect(connection.cancelRlmChild("finished-child")).resolves.toBe(false);
expect(fakeClient.requests[1]).toMatchObject({ type: "cancel_rlm_child", activeSessionId: "active-1", childId: "child-1", });
// A daemon from a build that predates the command reports a restart hint // instead of the raw protocol error. await expect(connection.cancelRlmChild("stale-daemon")).rejects.toThrow( "the daemon is running an older build; restart the daemon and try again", ); });
it("sends bash commands through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await connection.executeBash("echo hi", { excludeFromContext: true }); await connection.abortBash();
expect(fakeClient.requests[1]).toMatchObject({ type: "execute_bash", activeSessionId: "active-1", command: "echo hi", excludeFromContext: true, }); expect(fakeClient.requests[2]).toMatchObject({ type: "abort_bash", activeSessionId: "active-1" });
// A daemon from a build that predates the command reports a restart hint // instead of the raw protocol error. await expect(connection.executeBash("stale-daemon")).rejects.toThrow( "the daemon is running an older build; restart the daemon and try again", ); fakeClient.abortBashUnknownCommand = true; await expect(connection.abortBash()).rejects.toThrow( "the daemon is running an older build; restart the daemon and try again", ); });
it("loads resource snapshots through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect(connection.getResourceSnapshot()).resolves.toMatchObject({ contextFiles: [{ path: "/tmp/AGENTS.md" }], skills: [{ name: "demo-skill", filePath: "/tmp/skills/demo-skill/SKILL.md" }], diagnostics: { skills: [], prompts: [], extensions: [], themes: [] }, });
expect(fakeClient.requests[1]).toMatchObject({ type: "get_resource_snapshot", activeSessionId: "active-1", }); });
it("capability-gates ACP MCP server replacement", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "session-1"); expect(connection.supportsAcpMcpServers()).toBe(false); await expect( connection.replaceAcpMcpServers( [ { name: "task", type: "http", url: "https://task.example/mcp", headers: { Authorization: "Bearer task" }, }, ], "owner-a", ), ).rejects.toBeInstanceOf(DaemonCapabilityUnavailableError);
fakeClient.serverCapabilities.add("acp_mcp_servers"); expect(connection.supportsAcpMcpServers()).toBe(true); await connection.releaseAcpMcpServers("owner-a", ["task"]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "replace_acp_mcp_servers", ownerId: "owner-a", servers: [], }); await connection.dispose(); });
it("loads the full model catalog through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("model_catalog"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
const catalog = await connection.getModelCatalog();
expect(catalog.configuredProviders).toEqual(["openai"]); expect(catalog.models[0]).toMatchObject({ provider: "openai", id: "gpt-5.1" }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_model_catalog", activeSessionId: "active-1", }); });
it("falls back to configured models when the daemon lacks model catalog support", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
const catalog = await connection.getModelCatalog();
expect(catalog.configuredProviders).toEqual(["openai"]); expect(catalog.models[0]).toMatchObject({ provider: "openai", id: "gpt-5.1" }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_available_models", activeSessionId: "active-1", }); });
it("loads session context through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] }, }, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, emittedAt: "2026-01-01T00:00:00.000Z", }, });
await expect(connection.getSessionContext()).resolves.toEqual({ messages: [{ role: "user", content: "context prompt", timestamp: 3 }], thinkingLevel: "medium", model: { provider: "anthropic", modelId: "claude-sonnet-4-5" }, });
expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_context", activeSessionId: "active-1", }); });
it("loads session trees through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 0, steering: [], followUps: [] }, }, meta: { id: "active-1:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, emittedAt: "2026-01-01T00:00:00.000Z", }, });
await expect(connection.getSessionTree()).resolves.toEqual({ tree: [ { entry: { type: "message", id: "user-1", parentId: null, timestamp: "2026-01-01T00:00:00.000Z", message: { role: "user", content: "hello", timestamp: 1 }, }, children: [], }, ], leafId: "user-1", });
expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_tree", activeSessionId: "active-1", }); });
it("loads serializable tool metadata through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect(connection.getToolDefinition("custom_tool")).resolves.toEqual({ name: "custom_tool", label: "custom_tool", description: "custom_tool description", promptSnippet: "custom_tool prompt", promptGuidelines: ["Use custom_tool"], parameters: { type: "object" }, renderShell: "self", });
expect(fakeClient.requests[1]).toMatchObject({ type: "get_tool_definition", activeSessionId: "active-1", name: "custom_tool", }); });
it("forwards daemon session events through the connection boundary", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach();
fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 2, steering: ["interrupt"], followUps: ["later"] }, }, meta: { id: "active-1:14", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 14, emittedAt: "2026-01-01T00:00:00.000Z", }, }); fakeClient.emitMessage({ type: "session_event", activeSessionId: "other", event: { type: "session_action_update", actions: { queuedCount: 1, steering: ["ignored"], followUps: [] }, }, });
expect(events).toEqual([ { type: "session_event", event: { type: "session_action_update", actions: { queuedCount: 2, steering: ["interrupt"], followUps: ["later"] }, }, }, ]);
await connection.attach(); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", clientId: expect.any(String), capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", sequence: 14, }, });
await connection.attach(); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", activeSessionId: "active-1", capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], resumeCursor: { activeSessionId: "active-1", generation: "generation-active-1", sequence: 14, }, }); });
it("ignores delayed events from a retired daemon generation", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const steeringUpdates: Array<readonly string[]> = []; connection.subscribe((event) => { if (event.type === "session_event" && event.event.type === "session_action_update") { steeringUpdates.push(event.event.actions.steering); } }); await connection.attach(); const emitQueue = (generation: string, sequence: number, steering: string) => { fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "session_action_update", actions: { queuedCount: 1, steering: [steering], followUps: [] } }, meta: { id: `${generation}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence, cursor: { generation, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); };
emitQueue("generation-new", 1, "new"); emitQueue("generation-active-1", 13, "old"); await vi.waitFor(() => expect(steeringUpdates).toEqual([["new"]]));
await connection.attach(); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "attach", resumeCursor: { generation: "generation-new", sequence: 1 }, }); });
it("forwards extension UI requests and responses", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach(); expect(fakeClient.requests[0]).toMatchObject({ type: "attach", activeSessionId: "active-1", supportsExtensionUi: true, clientId: expect.any(String), capabilities: ["attach_snapshot", "event_sequence", "extension_ui", "slim_attach", "chunked_snapshot"], });
fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "request-1", method: "confirm", payload: { title: "Confirm", message: "Proceed?" }, }); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "other", id: "request-2", method: "confirm", payload: { title: "Ignore", message: "Wrong session" }, });
expect(events).toEqual([ { type: "extension_ui_request", request: { id: "request-1", method: "confirm", payload: { title: "Confirm", message: "Proceed?" }, }, }, ]);
await connection.respondToExtensionUiRequest("request-1", { confirmed: true }); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "extension_ui_response", activeSessionId: "active-1", requestId: "request-1", response: { confirmed: true }, }); });
it("uses an extended timeout for refine requests through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect( connection.refine({ instructions: "remember this", rollbackId: "refine_previous" }), ).resolves.toMatchObject({ id: "refine_daemon", appliedEdits: [], });
expect(fakeClient.requests[1]).toMatchObject({ type: "refine", activeSessionId: "active-1", instructions: "remember this", rollbackId: "refine_previous", }); expect(fakeClient.requests[1]).not.toHaveProperty("global"); expect(fakeClient.requestTimeouts[0]).toBe(30000); expect(fakeClient.requestTimeouts[1]).toBe(DAEMON_REFINE_REQUEST_TIMEOUT_MS); });
it("lists and renames saved sessions through the daemon protocol", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); const progress: Array<[number, number]> = []; const discovered: AgentConnectionSavedSessionInfo[] = [];
const sessions = await connection.listSavedSessions("current", { onProgress: (loaded, total) => { progress.push([loaded, total]); }, onSession: (session) => { discovered.push(session); }, }); await connection.renameSavedSession("/tmp/session-a.jsonl", "Next name"); await expect(connection.deleteSavedSession("/tmp/session-a.jsonl")).resolves.toEqual({ ok: true, method: "trash", });
expect(sessions).toEqual([ { path: "/tmp/session-a.jsonl", id: "session-a", cwd: "/tmp", name: "Saved session", created: new Date("2026-01-01T00:00:00.000Z"), modified: new Date("2026-01-02T00:00:00.000Z"), messageCount: 2, firstMessage: "hello", allMessagesText: "hello world", }, ]); expect(progress).toEqual([ [1, 2], [2, 2], ]); expect(discovered).toEqual(sessions); expect(fakeClient.requests[1]).toMatchObject({ type: "list_saved_sessions", activeSessionId: "active-1", scope: "current", }); expect(fakeClient.requests[2]).toMatchObject({ type: "rename_saved_session", activeSessionId: "active-1", sessionPath: "/tmp/session-a.jsonl", name: "Next name", }); expect(fakeClient.requests[3]).toMatchObject({ type: "delete_saved_session", activeSessionId: "active-1", sessionPath: "/tmp/session-a.jsonl", }); });
it("rehydrates daemon recoverable errors for interactive recovery flows", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
await expect(connection.switchSession("/tmp/session.jsonl")).rejects.toBeInstanceOf(MissingSessionCwdError); await expect(connection.switchSession("/tmp/session.jsonl")).rejects.toMatchObject({ issue: { sessionFile: "/tmp/session.jsonl", sessionCwd: "/tmp/missing", fallbackCwd: "/tmp/current", }, }); await expect(connection.importFromJsonl("/tmp/not-found.jsonl")).rejects.toBeInstanceOf( SessionImportFileNotFoundError, ); await expect(connection.importFromJsonl("/tmp/not-found.jsonl")).rejects.toMatchObject({ filePath: "/tmp/not-found.jsonl", }); });
it("can close the daemon client when disposed by an owning caller", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { closeClientOnDispose: true, }); await connection.attach();
await connection.dispose();
expect(fakeClient.requests.at(-1)).toMatchObject({ type: "detach", activeSessionId: "active-1" }); expect(fakeClient.closeCount).toBe(1); });
it("cleans up daemon client subscriptions when static attach fails", async () => { const fakeClient = new FakeDaemonClient();
await expect( DaemonAgentConnection.attach(asDaemonClient(fakeClient), "missing", { closeClientOnDispose: true }), ).rejects.toThrow("Unknown active session: missing");
expect(fakeClient.getMessageListenerCount()).toBe(0); expect(fakeClient.getCloseListenerCount()).toBe(0); expect(fakeClient.requests.map((request) => request.type)).toEqual(["attach", "detach"]); expect(fakeClient.closeCount).toBe(1); });});