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.
166 kB · 4605 lines
TypeScript
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799380038013802380338043805380638073808380938103811381238133814381538163817381838193820382138223823382438253826382738283829383038313832383338343835383638373838383938403841384238433844384538463847384838493850385138523853385438553856385738583859386038613862386338643865386638673868386938703871387238733874387538763877387838793880388138823883388438853886388738883889389038913892389338943895389638973898389939003901390239033904390539063907390839093910391139123913391439153916391739183919392039213922392339243925392639273928392939303931393239333934393539363937393839393940394139423943394439453946394739483949395039513952395339543955395639573958395939603961396239633964396539663967396839693970397139723973397439753976397739783979398039813982398339843985398639873988398939903991399239933994399539963997399839994000400140024003400440054006400740084009401040114012401340144015401640174018401940204021402240234024402540264027402840294030403140324033403440354036403740384039404040414042404340444045404640474048404940504051405240534054405540564057405840594060406140624063406440654066406740684069407040714072407340744075407640774078407940804081408240834084408540864087408840894090409140924093409440954096409740984099410041014102410341044105410641074108410941104111411241134114411541164117411841194120412141224123412441254126412741284129413041314132413341344135413641374138413941404141414241434144414541464147414841494150415141524153415441554156415741584159416041614162416341644165416641674168416941704171417241734174417541764177417841794180418141824183418441854186418741884189419041914192419341944195419641974198419942004201420242034204420542064207420842094210421142124213421442154216421742184219422042214222422342244225422642274228422942304231423242334234423542364237423842394240424142424243424442454246424742484249425042514252425342544255425642574258425942604261426242634264426542664267426842694270427142724273427442754276427742784279428042814282428342844285428642874288428942904291429242934294429542964297429842994300430143024303430443054306430743084309431043114312431343144315431643174318431943204321432243234324432543264327432843294330433143324333433443354336433743384339434043414342434343444345434643474348434943504351435243534354435543564357435843594360436143624363436443654366436743684369437043714372437343744375437643774378437943804381438243834384438543864387438843894390439143924393439443954396439743984399440044014402440344044405440644074408440944104411441244134414441544164417441844194420442144224423442444254426442744284429443044314432443344344435443644374438443944404441444244434444444544464447444844494450445144524453445444554456445744584459446044614462446344644465446644674468446944704471447244734474447544764477447844794480448144824483448444854486448744884489449044914492449344944495449644974498449945004501450245034504450545064507450845094510451145124513451445154516451745184519452045214522452345244525452645274528452945304531453245334534453545364537453845394540454145424543454445454546454745484549455045514552455345544555455645574558455945604561456245634564456545664567456845694570457145724573457445754576457745784579458045814582458345844585458645874588458945904591459245934594459545964597459845994600460146024603460446054606import { mkdtempSync, rmSync, writeFileSync } from "node:fs";import { tmpdir } from "node:os";import { join } from "node:path";import { PassThrough } from "node:stream";import type { AgentMessage } from "@earendil-works/pi-agent-core";import { fauxAssistantMessage, getModel } from "@earendil-works/pi-ai";import { describe, expect, it, vi } from "vitest";import { createManualContinuationMessage, MANUAL_CONTINUATION_PROMPT } from "../src/core/messages.js";import { DaemonAgentConnection } from "../src/modes/agent-connection/daemon-agent-connection.js";import type { AgentConnectionEvent, AgentConnectionRlmChildAgentSnapshot, AgentConnectionSessionTreeEntry, 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 { DaemonSessionStartingError } from "../src/modes/daemon/daemon-errors.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";import { InteractiveMode } from "../src/modes/interactive/interactive-mode.js";import { attachJsonlLineReader, serializeJsonLine } from "../src/modes/rpc/jsonl.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; attachError: Error | undefined; 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>> = []; createdActiveSessionId = "active-created"; createError: Error | undefined; createGate: Promise<void> | undefined; 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 "create": { if (this.createGate) await this.createGate; if (this.createError) throw this.createError; return { type: "response", command: command.type, success: true, data: { id: this.createdActiveSessionId, activeSessionId: this.createdActiveSessionId }, }; } case "attach": { if (this.attachFailures > 0) { this.attachFailures--; throw this.attachError ?? 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", }; } const response: DaemonResponse = { type: "response", command: command.type, success: true, data: this.attachResultFactory?.(command) ?? createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12), }; options.onResponse?.(response); return response; } 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_commands": return { type: "response", command: command.type, success: true, data: { commands: [ { name: "route", description: "Select a routing policy", argumentCompletions: [{ value: "policy ", label: "policy", description: "Follow a policy" }], source: "extension", sourceInfo: { path: "/tmp/extensions/route.ts", source: "local", scope: "user", origin: "top-level", }, }, ], }, }; 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: command.projection === "tree" ? { entries: [{ entry: projectedUserEntry("user-1", null) }], leafId: "user-1" } : { 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, }); const 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", }, ]; return { type: "response", command: command.type, success: true, data: { sessions: command.streamOnly === true ? [] : sessions, }, }; } 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;}
/** The projected tree entry the daemon sends for a user message. */function projectedUserEntry(id: string, parentId: string | null): AgentConnectionSessionTreeEntry { return { type: "message", id, parentId, timestamp: "2026-01-01T00:00:00.000Z", role: "user", text: "hello", hasText: true, };}
/** * Push a snapshot-staling session event so the next read goes back to the * daemon instead of the cached attach snapshot. */function emitStaleSessionActionUpdate(fakeClient: FakeDaemonClient): void { 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", }, });}
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([expect.objectContaining({ 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("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("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("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.each([false, true])("reattaches after an update restart (deferring=%s)", async (deferSessionEvents) => { 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", { deferSessionEvents, }); const events: AgentConnectionEvent[] = []; const restored = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { events.push(event); if (event.type === "session_resynced") { resolve(event); } }); }); // The reconnect settles without a timer: wait on the status sequence itself. let sawReconnecting = false; const connectedAgain = new Promise<void>((resolve) => { connection.subscribe((event) => { if (event.type !== "connection_status") return; if (event.status === "reconnecting") { sawReconnecting = true; } else if (sawReconnecting && event.status === "connected") { resolve(); } }); }); await connection.attach(); if (deferSessionEvents) emitSequencedSessionEvent(fakeClient, "active-original", 100);
fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-original", reason: "update", }); await restored; 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, }, }); await connection.flushBufferedSessionEvents(); 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 connectedAgain; 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("re-attaches and accepts prompt when connection drops with recoverDaemon", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); await connection.attach();
fakeClient.connected = false; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "shutdown"));
await connection.prompt("hello after restart");
expect(events.some((e) => e.type === "connection_status" && e.status === "reconnecting")).toBe(true); expect(events.some((e) => e.type === "session_resynced")).toBe(true); expect(events.some((e) => e.type === "connection_status" && e.status === "connected")).toBe(true); expect( fakeClient.requests.some( (r) => r.type === "prompt" && "message" in r && (r as { message: unknown }).message === "hello after restart", ), ).toBe(true); });
it("retries and succeeds when attach is refused with Session worker is starting", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const resynced = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") resolve(event); }); }); await connection.attach();
fakeClient.attachFailures = 1; fakeClient.attachError = new DaemonSessionStartingError(undefined, "Session worker is starting"); fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed"));
const event = await resynced; expect(event.type).toBe("session_resynced"); expect(fakeClient.resetTransportCount).toBe(0); });
it("emits a clear give-up error with transcript details after reconnect timeout expires", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("connection refused"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 50, }); const closed = new Promise<string>((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach();
fakeClient.connected = false; fakeClient.emitClose(new Error("socket died"));
const error = await closed; expect(error).toContain("Daemon reconnection failed"); expect(error).toContain("restart Prime Agent or reopen the session from Agents View"); expect(error).toContain("Session ID: session-current."); });
it("retries initial attach when worker is starting", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.attachFailures = 1; fakeClient.attachError = new DaemonSessionStartingError(undefined, "Session worker is starting"); const connection = await DaemonAgentConnection.attach(asDaemonClient(fakeClient), "active-original", { reconnectTimeoutMs: 5000, }); expect(connection).toBeDefined(); await connection.dispose(); });
it("waits for queued restore and remaps activeSessionId through list", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); const resynced = new Promise<AgentConnectionEvent>((resolve) => { connection.subscribe((event) => { if (event.type === "session_resynced") resolve(event); }); }); await connection.attach();
let listQueries = 0; const originalRequest = fakeClient.request.bind(fakeClient); fakeClient.request = async (command, ...args) => { if ( command.type === "attach" && "activeSessionId" in command && (command as any).activeSessionId === "active-original" ) { return { type: "response", command: command.type, success: false, error: "Unknown active session: active-original", }; } if (command.type === "list") { listQueries++; if (listQueries === 1) { return { type: "response", command: command.type, success: true, data: { sessions: [{ sessionId: "session-current", sessionFile: "/tmp/session.jsonl" }] }, }; } return { type: "response", command: command.type, success: true, data: { sessions: [ { activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session.jsonl", }, ], }, }; } return originalRequest(command, ...args); }; fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed"));
const event = await resynced; expect(event.type).toBe("session_resynced"); expect(listQueries).toBeGreaterThanOrEqual(2); expect((event as any).snapshot.state.activeSessionId).toBe("active-restored"); });
it("emits accurate update recovery error after update recovery timeout", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.reconnectError = new Error("connection refused"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 10, }); const closed = new Promise<string>((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach();
fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update"));
const error = await closed; expect(error).toContain("restarted for an update, but this window could not reconnect"); expect(error).not.toContain("daemon shut down while this window was attached"); expect(error).toContain("Session ID: session-current."); });
it.each([false, true])("bounds the reattach wait (saved transcript=%s)", async (transcriptSaved) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); const directory = mkdtempSync(join(tmpdir(), "prime-unrestorable-")); const transcript = join(directory, "session-current.jsonl"); if (transcriptSaved) writeFileSync(transcript, ""); try { const fakeClient = new FakeDaemonClient(); fakeClient.attachResultFactory = (command) => createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1, { state: { ...createConnectionState(command.activeSessionId, "session-current"), sessionFile: transcript }, }); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); const closed = new Promise<string>((resolve) => { connection.subscribe((event) => { if (event.type === "closed") resolve(event.error ?? ""); }); }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(300); // The bounded search, not the 120s update window, produced the verdict. const error = await closed; expect(error).toContain("is not in the restarted daemon's session list"); expect(error).toContain("Session ID: session-current."); expect(error).toContain(transcriptSaved ? "/resume session-current" : "start a fresh session with /new"); const readsAtVerdict = fakeClient.requests.filter((request) => request.type === "list").length; await vi.advanceTimersByTimeAsync(120_000); expect(fakeClient.requests.filter((request) => request.type === "list")).toHaveLength(readsAtVerdict); await expect(connection.prompt("after the restart")).rejects.toThrow("restarted daemon's session list"); } finally { rmSync(directory, { recursive: true, force: true }); vi.useRealTimers(); } }); it("does not carry a resolved roster miss into the next restart's first miss", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const events: AgentConnectionEvent[] = []; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); connection.subscribe((event) => void events.push(event)); await connection.attach();
// First restart: the roster misses the session for a moment, then lists it again. fakeClient.updateRestartSessions = []; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(20); expect(events.filter((event) => event.type === "closed")).toHaveLength(0); fakeClient.updateRestartSessions = [ { id: "active-original", sessionId: "session-current", activeSessionId: "active-restored" }, ]; await vi.advanceTimersByTimeAsync(120); expect(events.filter((event) => event.type === "closed")).toHaveLength(0);
// Much later, a second restart misses again: the fresh miss must wait its own bound. await vi.advanceTimersByTimeAsync(120_000); fakeClient.updateRestartSessions = []; fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(5); expect(events.filter((event) => event.type === "closed")).toHaveLength(0); await vi.advanceTimersByTimeAsync(400); expect(events.filter((event) => event.type === "closed")).toHaveLength(1); } finally { vi.useRealTimers(); } });
it("keeps waiting for a listed session and re-attaches it once its worker is routed", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.updateRestartSessions = [{ id: "active-original", sessionId: "session-current" }]; const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { sessionSearchTimeoutMs: 50, }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); // The daemon lists the session while its worker is still queued: the bound must not fire. await vi.advanceTimersByTimeAsync(600); fakeClient.updateRestartSessions = [ { id: "active-original", sessionId: "session-current", activeSessionId: "active-restored" }, ]; await vi.advanceTimersByTimeAsync(600); const attaches = fakeClient.requests.filter((request) => request.type === "attach"); expect(attaches.at(-1)).toMatchObject({ activeSessionId: "active-restored" }); await expect(connection.prompt("after the restore")).resolves.toBeUndefined(); } finally { vi.useRealTimers(); } });
async function attachLostConnection(fakeClient: FakeDaemonClient): Promise<{ connection: DaemonAgentConnection; events: AgentConnectionEvent[]; }> { const events: AgentConnectionEvent[] = []; const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); connection.subscribe((event) => void events.push(event)); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); await vi.advanceTimersByTimeAsync(300); return { connection, events }; }
it("starts a fresh session when /new is requested after the unrestorable verdict", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection, events } = await attachLostConnection(fakeClient);
await expect(connection.newSession()).resolves.toEqual({ cancelled: false }); expect(fakeClient.requests.find((request) => request.type === "create")).toMatchObject({ config: { cwd: "/tmp/project" }, }); expect(fakeClient.requests.filter((request) => request.type === "attach").at(-1)).toMatchObject({ activeSessionId: "active-fresh", }); expect(events.filter((event) => event.type === "session_replaced").at(-1)).toMatchObject({ state: { activeSessionId: "active-fresh" }, }); await expect(connection.prompt("after the recovery")).resolves.toBeUndefined(); } finally { vi.useRealTimers(); } });
it.each([undefined, "/tmp/elsewhere"])( "resumes a transcript after the verdict (cwdOverride=%s)", async (cwdOverride) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const { connection } = await attachLostConnection(fakeClient); await expect(connection.switchSession("/tmp/saved.jsonl", { cwdOverride })).resolves.toEqual({ cancelled: false, }); const create = fakeClient.requests.find((request) => request.type === "create"); expect(create).toMatchObject({ sessionPath: "/tmp/saved.jsonl" }); expect(create?.type === "create" ? create.config?.cwd : undefined).toBe(cwdOverride); } finally { vi.useRealTimers(); } }, );
it("reports an actionable error and stays retryable when the replacement cannot be created", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, sessionSearchTimeoutMs: 50, }); await connection.attach(); // This verdict arrives after a socket loss, which must leave the transport revivable. fakeClient.attachFailures = 99; fakeClient.attachError = new Error("Unknown active session: active-original"); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock")); await vi.advanceTimersByTimeAsync(2000); expect(fakeClient.closeCount).toBe(0); fakeClient.attachFailures = 0; fakeClient.createError = new Error("connect ECONNREFUSED"); await expect(connection.newSession()).rejects.toThrow(/ECONNREFUSED.*relaunch this window/); // The verdict survives: a retry reaches the daemon again instead of a poisoned connection. await expect(connection.newSession()).rejects.toThrow("relaunch this window"); expect(fakeClient.requests.filter((request) => request.type === "create")).toHaveLength(2); await expect(connection.prompt("x")).rejects.toThrow("restarted daemon's session list"); } finally { vi.useRealTimers(); } });
it("reports a failed session command instead of exiting the client after the terminal close", async () => { const shown: string[] = []; const listeners = new Set<(event: AgentConnectionEvent) => void>(); const mode = Object.create(InteractiveMode.prototype) as { connectionClosed: boolean; handleClearCommand(): Promise<void>; subscribeToAgent(): void; }; Object.assign(mode, { stop: () => {}, stopWorkingLoader: () => {}, clearStatus: () => {}, agentConnection: { subscribe: (listener: (event: AgentConnectionEvent) => void) => { listeners.add(listener); return () => listeners.delete(listener); }, newSession: async () => { throw new Error("The daemon could not start a replacement session: ECONNREFUSED. Relaunch this window."); }, }, showError: (message: string) => { shown.push(message); }, }); const exit = vi.spyOn(process, "exit").mockImplementation((code) => { throw new Error(`process.exit(${String(code)})`); }); mode.subscribeToAgent(); for (const listener of listeners) await listener({ type: "closed", error: "restarted daemon's session list" }); expect(mode.connectionClosed).toBe(true); await expect(mode.handleClearCommand()).resolves.toBeUndefined(); expect(exit).not.toHaveBeenCalled(); expect(shown.join(" ")).toContain("Relaunch this window"); });
it("remaps activeSessionId and succeeds on request failing with Unknown active session", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 5000, }); await connection.attach();
fakeClient.request = async (command) => { if ( "activeSessionId" in command && (command as { activeSessionId: unknown }).activeSessionId === "active-original" ) { return { type: "response", command: command.type, success: false, error: "Unknown active session: active-original", }; } if (command.type === "list") { return { type: "response", command: command.type, success: true, data: { sessions: [ { activeSessionId: "active-restored", sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", }, ], }, }; } if (command.type === "get_session_stats") { return { type: "response", command: command.type, success: true, data: { sessionId: "session-current", sessionFile: "/tmp/session-current.jsonl", userMessages: 1, assistantMessages: 1, toolCalls: 0, toolResults: 0, totalMessages: 2, }, }; } return { type: "response", command: command.type, success: true, data: createAttachResult("active-restored", "fake-client", [], 12), }; };
const stats = await connection.getSessionStats(); expect(stats.sessionId).toBe("session-current"); });
it("survives a request when remap fails without unhandled crash", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-original", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 10, }); await connection.attach();
fakeClient.request = async (command) => ({ type: "response", command: command.type, success: false, error: "Unknown active session: active-original", });
await expect(connection.getSessionStats()).rejects.toThrow(); }); it("bounds remap wait on unlisted session and recovers via /new", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.promptResponseError = "Unknown active session: active-original"; const events: AgentConnectionEvent[] = []; const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); connection.subscribe((e) => void events.push(e)); await connection.attach(); fakeClient.attachFailures = 999; fakeClient.attachError = new Error("Unknown active session: active-original"); const promptRejection = expect(connection.prompt("hi")).rejects.toThrow("restarted daemon's session list"); await vi.advanceTimersByTimeAsync(300); await promptRejection; expect(events.filter((e) => e.type === "closed")).toHaveLength(1); fakeClient.attachFailures = 0; fakeClient.createdActiveSessionId = "active-fresh"; await expect(connection.newSession()).resolves.toEqual({ cancelled: false }); expect(events.filter((e) => e.type === "session_replaced")).toHaveLength(1); } finally { vi.useRealTimers(); } }); it("single-flights concurrent /new in the verdict state", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection } = await attachLostConnection(fakeClient); const results = await Promise.all([connection.newSession(), connection.newSession()]); expect(results).toEqual([{ cancelled: false }, { cancelled: false }]); expect(fakeClient.requests.filter((r) => r.type === "create")).toHaveLength(1); expect(fakeClient.requests.filter((r) => r.type === "new_session")).toHaveLength(0); } finally { vi.useRealTimers(); } }); it("releases created session when disposed mid-replacement without unhandled error", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); fakeClient.createdActiveSessionId = "active-fresh"; const { connection } = await attachLostConnection(fakeClient); let release!: () => void; fakeClient.createGate = new Promise((r) => { release = r; }); const [replace, dispose] = [connection.newSession(), connection.dispose()]; release(); await Promise.all([replace, dispose]); const releases = fakeClient.requests.filter((r) => r.type === "detach" || r.type === "complete_owned_session"); expect(releases.at(-1)).toMatchObject({ activeSessionId: "active-fresh" }); } finally { vi.useRealTimers(); } }); it("rejects queued prompt with unrestorable verdict when reconnect produces verdict", async () => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); try { const fakeClient = new FakeDaemonClient(); const client = asDaemonClient(fakeClient); const connection = new DaemonAgentConnection(client, "active-original", { sessionSearchTimeoutMs: 50 }); await connection.attach(); fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); const promptRejection = expect(connection.prompt("queued")).rejects.toThrow("restarted daemon's session list"); await vi.advanceTimersByTimeAsync(300); await promptRejection; } finally { vi.useRealTimers(); } });
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("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: projectedUserEntry("user-1", null), 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("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.each([ ["headless", "message_end"], ["headless", "message_update"], ["reconnect", "message_end"], ["reconnect", "message_update"], ] as const)("preserves %s %s coalesced with attach snapshot completion", async (mode, eventType) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: mode === "reconnect", recoverDaemon: async () => {}, }); const input = new PassThrough(); const detachReader = attachJsonlLineReader(input, (line) => fakeClient.emitMessage(JSON.parse(line))); try { if (mode === "reconnect") { await connection.attach(); await connection.flushBufferedSessionEvents(); } const message = fauxAssistantMessage("new response"); fakeClient.attachResultFactory = (command) => { const full = createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12, { state: { ...createConnectionState("active-1", "session-current"), isStreaming: true }, streamingMessage: fauxAssistantMessage("old partial response"), }); const { messages: _messages, ...snapshot } = full.snapshot; const records: DaemonOutbound[] = [ { type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "coalesced", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, }, { type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "coalesced", chunkCount: 0, lastEventSequence: 12, lastEventCursor: full.lastEventCursor, }, { type: "session_event", activeSessionId: "active-1", event: eventType === "message_end" ? { type: eventType, message } : { type: eventType, message, assistantMessageEvent: { type: "text_delta", contentIndex: 0, delta: "response", partial: message, }, }, 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", }, }, ]; // All records dispatch before the attach response's promise continuation. queueMicrotask(() => input.write(records.map(serializeJsonLine).join(""))); return { ...full, snapshotStream: { id: "coalesced", messageCount: 0, targetChunkBytes: 512 * 1024 } }; }; const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); if (mode === "reconnect") { fakeClient.connected = false; fakeClient.emitClose(new Error("socket closed")); await vi.waitFor(() => expect(events.some((event) => event.type === "session_resynced")).toBe(true)); } else { await connection.attach(); } const snapshot = await connection.getInitialSnapshot(); expect(snapshot.streamingMessage).toEqual(eventType === "message_end" ? undefined : message); expect(snapshot.lastEventCursor).toEqual({ generation: "generation-active-1", sequence: 13 }); expect(fakeClient.requests.map((request) => request.type)).toContain("get_messages"); expect(events).toContainEqual( expect.objectContaining({ type: "session_event", event: expect.objectContaining({ type: eventType }) }), ); for (const event of events) { if (event.type === "session_resynced") expect(event.snapshot.streamingMessage).toEqual(snapshot.streamingMessage); } } finally { detachReader(); input.destroy(); await connection.dispose(); } });
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((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.each([ ["reconnect", "attach"], ["reconnect", "snapshot"], ["reconnect", "backoff"], ["update", "attach"], ["update", "snapshot"], ["update", "backoff"], ] as const)("stops %s after a terminal close during %s", async (recovery, stage) => { vi.useFakeTimers({ toFake: ["setTimeout", "clearTimeout", "Date"] }); const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { recoverDaemon: async () => undefined, reconnectTimeoutMs: 1000, }); let release!: () => void; const gate = new Promise<void>((resolve) => { release = resolve; }); try { await connection.attach(); fakeClient.updateRestartSessions = [{ activeSessionId: "active-1", sessionId: "session-current" }]; const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (...args) => { const response = await request(...args); if (args[0].type === "attach") { if (stage === "backoff") throw new Error("attach temporarily unavailable"); if (stage === "attach") await gate; } return response; }); const getInitialSnapshot = connection.getInitialSnapshot.bind(connection); const snapshotRead = vi.spyOn(connection, "getInitialSnapshot").mockImplementation(async (...args) => { const snapshot = await getInitialSnapshot(...args); if (stage === "snapshot") await gate; return snapshot; }); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); fakeClient.connected = false; fakeClient.emitClose( recovery === "update" ? new DaemonSocketClosedError("/tmp/fake.sock", "update") : new Error("Daemon socket closed"), ); await vi.advanceTimersByTimeAsync(0); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(2); if (stage === "snapshot") expect(snapshotRead).toHaveBeenCalledOnce(); const requestCount = fakeClient.requests.length; const reconnectCount = fakeClient.reconnectCount; const resetCount = fakeClient.resetTransportCount; const closeCount = fakeClient.closeCount; fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); release(); await vi.advanceTimersByTimeAsync(120100); expect(fakeClient.requests).toHaveLength(requestCount); expect(fakeClient.reconnectCount).toBe(reconnectCount); expect(fakeClient.resetTransportCount).toBe(resetCount); expect(fakeClient.closeCount).toBe(closeCount); expect(events).toEqual([ expect.objectContaining({ type: "connection_status", status: "reconnecting" }), expect.objectContaining({ type: "closed", error: expect.stringContaining("The daemon stopped this agent session."), }), ]); } finally { release(); await connection.dispose(); vi.useRealTimers(); } });
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("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("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("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("asks for the projected session tree when the daemon supports it", async () => { const fakeClient = new FakeDaemonClient(); fakeClient.serverCapabilities.add("session_tree_projection"); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitStaleSessionActionUpdate(fakeClient);
await expect(connection.getSessionTree()).resolves.toEqual({ tree: [{ entry: projectedUserEntry("user-1", null), children: [] }], leafId: "user-1", }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_tree", activeSessionId: "active-1", projection: "tree", }); });
it("projects the session tree locally when the daemon lacks the projection", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach(); emitStaleSessionActionUpdate(fakeClient);
await expect(connection.getSessionTree()).resolves.toEqual({ tree: [{ entry: projectedUserEntry("user-1", null), children: [] }], leafId: "user-1", }); expect(fakeClient.requests[1]).toMatchObject({ type: "get_session_tree", activeSessionId: "active-1" }); expect(fakeClient.requests[1]).not.toHaveProperty("projection"); });
it("capability-gates published extension command argument completions", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); await connection.attach();
// An older daemon may still send the field; the client must not complete from it. const withoutCapability = await connection.getCommands(); expect(withoutCapability[0]?.name).toBe("route"); expect(withoutCapability[0]?.argumentCompletions).toBeUndefined();
fakeClient.serverCapabilities.add("command_argument_completions"); const withCapability = await connection.getCommands(); expect(withCapability[0]?.argumentCompletions).toEqual([ { value: "policy ", label: "policy", description: "Follow a policy" }, ]); expect(fakeClient.requests.at(-1)).toMatchObject({ type: "get_commands", activeSessionId: "active-1" }); });
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 }, }); });});
const DEFERRAL_EVENT_BASE_SEQUENCE = 13;
function emitSequencedSessionEvent(client: FakeDaemonClient, activeSessionId: string, sequence: number): void { client.emitMessage({ type: "session_event", activeSessionId, event: { type: "session_info_changed", name: String(sequence) }, meta: { id: `${activeSessionId}:${sequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId, sequence, cursor: { generation: `generation-${activeSessionId}`, sequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, } satisfies DaemonOutbound);}
async function nextMessageLoopTurn(): Promise<void> { await new Promise<void>((resolve) => setImmediate(resolve)); await new Promise<void>((resolve) => setImmediate(resolve));}
describe("DaemonAgentConnection deferred session events", () => { it("defers events between attach and flush so the attach snapshot stays usable", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); });
emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); await nextMessageLoopTurn(); expect(delivered).toEqual([]);
fakeClient.requests.length = 0; const snapshot = await connection.getInitialSnapshot(); expect(fakeClient.requests.map((request) => request.type)).not.toContain("get_messages"); expect(snapshot.state.activeSessionId).toBe("active-1");
await connection.flushBufferedSessionEvents(); expect(delivered).toHaveLength(1); expect(delivered[0]).toMatchObject({ type: "session_event" });
emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + 1); await nextMessageLoopTurn(); expect(delivered).toHaveLength(2); } finally { await connection.dispose(); } });
it("delivers events live and refetches when deferral is not opted in", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); });
emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); await nextMessageLoopTurn(); expect(delivered).toHaveLength(1);
fakeClient.requests.length = 0; await connection.getInitialSnapshot(); expect(fakeClient.requests.map((request) => request.type)).toContain("get_messages"); } finally { await connection.dispose(); } });
it("drops deferred events a resync snapshot already contains", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); });
const resyncSequence = DEFERRAL_EVENT_BASE_SEQUENCE + 2; emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE); emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + 1); fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot: createAttachResult("active-1", undefined, undefined, resyncSequence).snapshot, meta: { id: `active-1:${resyncSequence}`, protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: resyncSequence, cursor: { generation: "generation-active-1", sequence: resyncSequence }, emittedAt: "2026-01-01T00:00:00.000Z", }, } satisfies DaemonOutbound); await nextMessageLoopTurn(); expect(delivered.filter((event) => event.type === "session_resynced")).toHaveLength(1); expect(delivered.filter((event) => event.type === "session_event")).toEqual([]);
await connection.flushBufferedSessionEvents(); expect(delivered.filter((event) => event.type === "session_event")).toEqual([]);
emitSequencedSessionEvent(fakeClient, "active-1", resyncSequence + 1); await connection.flushBufferedSessionEvents(); expect(delivered.filter((event) => event.type === "session_event")).toHaveLength(1); } finally { await connection.dispose(); } });
it("resyncs an already rendered snapshot when deferral overflows", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const renderedSnapshot = await connection.getInitialSnapshot(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); });
// The initial renderer already has its snapshot when the buffer overflows. for (let index = 0; index <= 1000; index++) { emitSequencedSessionEvent(fakeClient, "active-1", DEFERRAL_EVENT_BASE_SEQUENCE + index); } await nextMessageLoopTurn(); expect(delivered).toHaveLength(0); const recoveredMessages: AgentMessage[] = [{ role: "user", content: "recovered", timestamp: 1 }]; fakeClient.attachResultFactory = (command) => { emitSequencedSessionEvent(fakeClient, "active-1", 1014); return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 1013, { messages: recoveredMessages, }); }; await connection.flushBufferedSessionEvents(); expect(renderedSnapshot.messages).toEqual([]); expect(delivered).toEqual([ expect.objectContaining({ type: "session_resynced", snapshot: expect.objectContaining({ messages: recoveredMessages }), }), { type: "session_event", event: { type: "session_info_changed", name: "1014" } }, ]);
emitSequencedSessionEvent(fakeClient, "active-1", 1015); await nextMessageLoopTurn(); expect(delivered).toHaveLength(3); } finally { await connection.dispose(); } });
it.each([true, false])( "finishes overflow recovery under continuous traffic (snapshot covers overflow=%s)", async (covered) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const rendering = new Promise<void>((resolve) => { release = resolve; }); try { await connection.attach(); let sequence = 12; for (let index = 0; index <= 1000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); let attaches = 0; fakeClient.attachResultFactory = (command) => { if (++attaches > 3) throw new Error("overflow recovery did not terminate"); const before = sequence; if (covered || attaches === 1) { for (let index = 0; index <= 1000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); } return createAttachResult( command.activeSessionId, command.clientId, command.capabilities, !covered && attaches === 1 ? before : sequence, ); }; const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); if (event.type === "session_resynced") return rendering; }); const flush = connection.flushBufferedSessionEvents(); await nextMessageLoopTurn(); expect(attaches).toBe(covered ? 1 : 2); expect(delivered.at(-1)).toMatchObject({ type: "session_resynced", snapshot: { lastEventSequence: sequence }, }); const beforeLive = delivered.length; for (let index = 0; index < 2000; index++) emitSequencedSessionEvent(fakeClient, "active-1", ++sequence); expect(delivered.slice(beforeLive)).toHaveLength(2000); release(); await flush; expect(attaches).toBe(covered ? 1 : 2); expect(delivered.some((event) => event.type === "closed")).toBe(false); } finally { release(); await connection.dispose(); } }, );
it.each([ ["replacement", false], ["replacement", true], ["streamed replacement", true], ["reattach", false], ["reattach", true], ] as const)("discards overflow recovery across %s (recovery started=%s)", async (change, recoveryStarted) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const gate = new Promise<void>((resolve) => { release = resolve; }); try { await connection.attach(); for (let sequence = 13; sequence <= 1013; sequence++) emitSequencedSessionEvent(fakeClient, "active-1", sequence); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (command, ...options) => { if (command.type === "attach" && change === "streamed replacement") { const result = createAttachResult("active-1", undefined, undefined, 1013); result.snapshotStream = { id: "old-recovery", messageCount: 0, targetChunkBytes: 512 * 1024 }; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "old-recovery", snapshot: result.snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); return { type: "response", command: command.type, success: true, data: result }; } if (command.type === "attach") await gate; if (command.type === "switch_session") return { type: "response", command: command.type, success: false, error: "Session already active", errorInfo: { code: "session_already_active", activeSessionId: "active-2", sessionPath: "/tmp/target.jsonl", }, }; if (command.type === "reattach") return { type: "response", command: command.type, success: true, data: createAttachResult("active-2", undefined, undefined, 1), }; return request(command, ...options); }); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); const flush = recoveryStarted ? connection.flushBufferedSessionEvents() : undefined; await nextMessageLoopTurn(); if (change === "reattach") { await connection.switchSession("/tmp/target.jsonl"); } else { fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); } const activeSessionId = change === "reattach" ? "active-2" : "active-1"; const sequence = change === "reattach" ? 2 : 1014; emitSequencedSessionEvent(fakeClient, activeSessionId, sequence); if (change === "streamed replacement") { fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "old-recovery", chunkCount: 0, lastEventSequence: 1013, lastEventCursor: { generation: "generation-active-1", sequence: 1013 }, }); } release(); await (flush ?? connection.flushBufferedSessionEvents()); expect(delivered.map((event) => event.type)).toEqual(["session_replaced", "session_event"]); expect(delivered[1]).toEqual({ type: "session_event", event: { type: "session_info_changed", name: String(sequence) }, }); if (change !== "reattach") expect( (connection as unknown as { latestSnapshot: { state: AgentConnectionState } }).latestSnapshot.state .sessionId, ).toBe("replacement"); expect((await connection.getState()).activeSessionId).toBe(activeSessionId); } finally { release(); await connection.dispose(); } });
it("does not advance event cursors for an attach superseded before its snapshot finishes", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); fakeClient.attachResultFactory = (command) => ({ ...createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 100), snapshotStream: { id: "superseded", messageCount: 0, targetChunkBytes: 512 * 1024 }, }); const pendingAttach = connection.attach(); await nextMessageLoopTurn(); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "superseded", snapshot: createAttachResult("active-1", undefined, undefined, 100).snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "superseded", chunkCount: 0, lastEventSequence: 100, lastEventCursor: { generation: "generation-active-1", sequence: 100 }, }); await pendingAttach; const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); emitSequencedSessionEvent(fakeClient, "active-1", 13); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "session_info_changed", name: "13" } }]); const snapshot = await connection.getInitialSnapshot(); expect(snapshot.lastEventCursor).toEqual({ generation: "generation-active-1", sequence: 13 }); } finally { await connection.dispose(); } });
it.each([0, 2, 3])( "releases a superseded attach with %i snapshot frames received before its response", async (frameCount) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); let release!: () => void; const responseGate = new Promise<void>((resolve) => { release = resolve; }); try { await connection.attach(); const result = createAttachResult("active-1", undefined, undefined, 100); fakeClient.attachResultFactory = () => ({ ...result, snapshotStream: { id: "abandoned", messageCount: 1, targetChunkBytes: 512 * 1024 }, }); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (...args) => { const response = await request(...args); if (args[0].type === "attach") await responseGate; return response; }); const frames: DaemonOutbound[] = [ { type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "abandoned", snapshot: result.snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose: "attach", }, { type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId: "abandoned", index: 0, messages: [{ role: "user", content: "old transcript", timestamp: 1 }], }, { type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "abandoned", chunkCount: 1, lastEventSequence: 100, }, ]; const pendingAttach = connection.attach(); for (const frame of frames.slice(0, frameCount)) fakeClient.emitMessage(frame); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); release(); await pendingAttach; expect( (connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size, ).toBe(0); for (const frame of frames.slice(frameCount)) fakeClient.emitMessage(frame); await nextMessageLoopTurn(); expect( (connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size, ).toBe(0); expect((await connection.getInitialSnapshot()).state.sessionId).toBe("replacement"); } finally { release(); await connection.dispose(); } }, );
it.each(["before response", "during stream", "after end"] as const)( "rejects a closed attach immediately %s", async (timing) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); const result = createAttachResult("active-1", undefined, undefined, 12); fakeClient.attachResultFactory = () => ({ ...result, snapshotStream: { id: "closed", messageCount: 0, targetChunkBytes: 512 * 1024 }, }); const rejected = vi.fn(); const resolved = vi.fn(); const events: AgentConnectionEvent[] = []; connection.subscribe((event) => { events.push(event); }); const pendingAttach = connection.attach().then(resolved, rejected); try { if (timing !== "before response") await nextMessageLoopTurn(); fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "closed", snapshot: result.snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "attach", }); if (timing === "after end") { fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "closed", chunkCount: 0, lastEventSequence: 12, }); } fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); await nextMessageLoopTurn(); expect(rejected).toHaveBeenCalledExactlyOnceWith(expect.any(Error)); expect(resolved).not.toHaveBeenCalled(); expect( (connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size, ).toBe(0); for (const purpose of ["attach", "replacement", "resync"] as const) { fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: purpose, snapshot: result.snapshot, messageCount: 1, targetChunkBytes: 512 * 1024, purpose, }); fakeClient.emitMessage({ type: "session_snapshot_chunk", activeSessionId: "active-1", snapshotId: purpose, index: 0, messages: [{ role: "user", content: "late transcript", timestamp: 1 }], }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: purpose, chunkCount: 1, lastEventSequence: 12, }); } fakeClient.emitMessage({ type: "session_snapshot_failed", activeSessionId: "active-1", snapshotId: "failed-after-close", error: "stream closed", }); expect( (connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size, ).toBe(0); fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "killed" }); await nextMessageLoopTurn(); expect( (connection as unknown as { snapshotAssemblies: Map<string, unknown> }).snapshotAssemblies.size, ).toBe(0); expect(events).toEqual([expect.objectContaining({ type: "closed" })]); const requestCount = fakeClient.requests.length; await expect(connection.attach()).rejects.toThrow("closed"); expect(fakeClient.requests).toHaveLength(requestCount); } finally { await connection.dispose(); await pendingAttach; } }, );
it.each(["session_replaced", "session_resynced"] as const)( "drops a %s superseded during the initial render", async (type) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); let release!: () => void; const initialRenderPromise = new Promise<void>((resolve) => { release = resolve; }); const ui = { agentConnection: connection, sessionEventQueue: Promise.resolve(), sessionEventGeneration: 0, initialRenderPromise, refreshCommandCatalogForCurrentSession: vi.fn(async () => {}), renderResyncedSession: vi.fn(async () => {}), resetSideQuestion: vi.fn(), resetExtensionUI: vi.fn(), applyConnectionStateSnapshot: vi.fn(), resetCurrentSessionRenderState: vi.fn(), rebindCurrentSession: vi.fn(async () => {}), renderInitialMessages: vi.fn(async () => {}), ui: { requestRender: vi.fn() }, showError: vi.fn(), }; try { await connection.attach(); (InteractiveMode.prototype as unknown as { subscribeToAgent(this: typeof ui): void }).subscribeToAgent.call( ui, ); const snapshot = createAttachResult("active-1", undefined, undefined, 12).snapshot; fakeClient.emitMessage( type === "session_replaced" ? { type, activeSessionId: "active-1", state: snapshot.state, messages: [] } : { type, activeSessionId: "active-1", snapshot }, ); await nextMessageLoopTurn(); const state = createConnectionState("active-1", "latest"); fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state, messages: [] }); release(); await ui.sessionEventQueue; expect(ui.renderResyncedSession).not.toHaveBeenCalled(); expect(ui.resetSideQuestion).toHaveBeenCalledOnce(); expect(ui.resetExtensionUI).toHaveBeenCalledOnce(); expect(ui.applyConnectionStateSnapshot).toHaveBeenCalledExactlyOnceWith(state); expect(ui.resetCurrentSessionRenderState).toHaveBeenCalledOnce(); expect(ui.rebindCurrentSession).toHaveBeenCalledOnce(); expect(ui.renderInitialMessages).toHaveBeenCalledOnce(); expect(ui.ui.requestRender).toHaveBeenCalledOnce(); expect(ui.showError).not.toHaveBeenCalled(); } finally { release(); await connection.dispose(); } }, );
it("keeps live events behind the whole replay and shares concurrent flushes", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); let release!: () => void; const gate = new Promise<void>((resolve) => { release = resolve; }); try { await connection.attach(); const delivered: string[] = []; const ui = { agentConnection: connection, sessionEventQueue: Promise.resolve(), sessionEventGeneration: 0, handleEvent: async (event: { type: string; name?: string }) => { delivered.push(event.name!); if (delivered.length === 1) await gate; }, showError: vi.fn(), }; (InteractiveMode.prototype as unknown as { subscribeToAgent(this: typeof ui): void }).subscribeToAgent.call( ui, ); emitSequencedSessionEvent(fakeClient, "active-1", 13); emitSequencedSessionEvent(fakeClient, "active-1", 14); const flush = connection.flushBufferedSessionEvents(); await nextMessageLoopTurn(); emitSequencedSessionEvent(fakeClient, "active-1", 15); expect(connection.flushBufferedSessionEvents()).toBe(flush); expect(delivered).toEqual(["13"]); release(); await flush; await ui.sessionEventQueue; expect(delivered).toEqual(["13", "14", "15"]); expect(ui.showError).not.toHaveBeenCalled(); } finally { release(); await connection.dispose(); } });
it("leaves streaming state in the attach snapshot unchanged until replay", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); const message = fauxAssistantMessage("streaming response"); fakeClient.emitMessage({ type: "session_event", activeSessionId: "active-1", event: { type: "message_start", message }, }); const snapshot = await connection.getInitialSnapshot(); expect(snapshot.streamingMessage).toBeUndefined(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "message_start", message } }]); expect((await connection.getInitialSnapshot()).streamingMessage).toEqual(message); } finally { await connection.dispose(); } });
it.each([true, false])( "delivers attach-time extension prompts without skipping transcript events (subscribed=%s)", async (subscribed) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { const delivered: AgentConnectionEvent[] = []; const listener = (event: AgentConnectionEvent) => { delivered.push(event); }; if (subscribed) connection.subscribe(listener); fakeClient.attachResultFactory = (command) => { fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "editor-request", method: "editor", payload: { title: "Edit" }, meta: { id: "active-1:14", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 14, cursor: { generation: "generation-active-1", sequence: 14 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); return createAttachResult(command.activeSessionId, command.clientId, command.capabilities, 12); }; await connection.attach(); // The transport releases earlier transcript events only after the snapshot. emitSequencedSessionEvent(fakeClient, "active-1", 13); if (!subscribed) { expect(delivered).toEqual([]); connection.subscribe(listener); } expect(delivered).toEqual([ { type: "extension_ui_request", request: { id: "editor-request", method: "editor", payload: { title: "Edit" } }, }, ]); await connection.flushBufferedSessionEvents(); expect(delivered.at(-1)).toEqual({ type: "session_event", event: { type: "session_info_changed", name: "13" }, }); const additionalListener = vi.fn(); connection.subscribe(additionalListener); expect(additionalListener).not.toHaveBeenCalled(); } finally { await connection.dispose(); } }, );
it.each(["notify", "setStatus", "setWidget"])( "bounds attach-time %s updates without dropping queued dialogs", async (method) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "dialog", method: "editor", payload: { title: "Edit" }, }); for (let index = 0; index < 2000; index++) { fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: String(index), method, payload: { message: String(index), statusKey: "progress", statusText: String(index), widgetKey: "progress", widgetLines: [String(index)], }, }); } const listener = vi.fn(); connection.subscribe(listener); expect(listener).toHaveBeenCalledTimes(128); expect(listener).toHaveBeenNthCalledWith(1, { type: "extension_ui_request", request: { id: "dialog", method: "editor", payload: { title: "Edit" } }, }); expect(listener).toHaveBeenLastCalledWith({ type: "extension_ui_request", request: expect.objectContaining({ id: "1999", method }), }); expect(fakeClient.requests.filter((request) => request.type === "extension_ui_response")).toEqual([]); } finally { await connection.dispose(); } }, );
it.each(["disposed", "replaced", "restarted", "update-session", "update-transport"])( "discards queued extension prompts when %s before subscription", async (boundary) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1"); try { await connection.attach(); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "old", method: "editor", payload: { title: "Old" }, meta: { id: "old:13", protocol: DAEMON_PROTOCOL_INFO, activeSessionId: "active-1", sequence: 13, cursor: { generation: "generation-active-1", sequence: 13 }, emittedAt: "2026-01-01T00:00:00.000Z", }, }); if (boundary === "disposed") await connection.dispose(); else if (boundary === "replaced") fakeClient.emitMessage({ type: "session_replaced", activeSessionId: "active-1", state: createConnectionState("active-1", "replacement"), messages: [], }); else if (boundary.startsWith("update-")) { fakeClient.updateRestartSessions = [{ activeSessionId: "active-1", sessionId: "session-current" }]; if (boundary === "update-session") { fakeClient.emitMessage({ type: "session_closed", activeSessionId: "active-1", reason: "update" }); } else { fakeClient.emitClose(new DaemonSocketClosedError("/tmp/prime-agent.sock", "update")); } } else { const snapshot = createAttachResult("active-1", undefined, undefined, 1).snapshot; snapshot.lastEventCursor = { generation: "restarted", sequence: 1 }; fakeClient.emitMessage({ type: "session_resynced", activeSessionId: "active-1", snapshot }); } let resolveConnected!: () => void; const connected = new Promise<void>((resolve) => { resolveConnected = resolve; }); const listener = vi.fn((event: AgentConnectionEvent) => { if (event.type === "connection_status" && event.status === "connected") resolveConnected(); }); connection.subscribe(listener); expect(listener).not.toHaveBeenCalled(); if (boundary.startsWith("update-")) { await connected; expect(listener).toHaveBeenCalledWith({ type: "connection_status", status: "connected" }); expect(listener).not.toHaveBeenCalledWith(expect.objectContaining({ type: "extension_ui_request" })); fakeClient.emitMessage({ type: "extension_ui_request", activeSessionId: "active-1", id: "new", method: "editor", payload: { title: "New" }, }); expect(listener).toHaveBeenLastCalledWith({ type: "extension_ui_request", request: { id: "new", method: "editor", payload: { title: "New" } }, }); } } finally { await connection.dispose(); } }, );
it.each([1, 1001])( "drops %i buffered events and overflow superseded by a restarted worker snapshot", async (count) => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); for (let index = 0; index < count; index++) emitSequencedSessionEvent(fakeClient, "active-1", 100 + index); const snapshot = createAttachResult("active-1", undefined, undefined, 1).snapshot; snapshot.lastEventCursor = { generation: "restarted", sequence: 1 }; fakeClient.emitMessage({ type: "session_snapshot_begin", activeSessionId: "active-1", snapshotId: "restart", snapshot, messageCount: 0, targetChunkBytes: 512 * 1024, purpose: "resync", }); fakeClient.emitMessage({ type: "session_snapshot_end", activeSessionId: "active-1", snapshotId: "restart", chunkCount: 0, lastEventSequence: 1, lastEventCursor: snapshot.lastEventCursor, }); await nextMessageLoopTurn(); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([]); expect(fakeClient.requests.filter((request) => request.type === "attach")).toHaveLength(1); expect((await connection.getInitialSnapshot()).lastEventSequence).toBe(1); } finally { await connection.dispose(); } }, );
it("replays only newer target-session events when reattaching with an inline snapshot", async () => { const fakeClient = new FakeDaemonClient(); const connection = new DaemonAgentConnection(asDaemonClient(fakeClient), "active-1", { deferSessionEvents: true, }); try { await connection.attach(); emitSequencedSessionEvent(fakeClient, "active-1", 13); const request = fakeClient.request.bind(fakeClient); vi.spyOn(fakeClient, "request").mockImplementation(async (command, ...options) => { if (command.type === "switch_session") return { type: "response", command: command.type, success: false, error: "Session already active", errorInfo: { code: "session_already_active", activeSessionId: "active-2", sessionPath: "/tmp/target.jsonl", }, }; if (command.type === "reattach") { emitSequencedSessionEvent(fakeClient, "active-2", 2); return { type: "response", command: command.type, success: true, data: createAttachResult("active-2", undefined, undefined, 1), }; } return request(command, ...options); }); await connection.switchSession("/tmp/target.jsonl"); const delivered: AgentConnectionEvent[] = []; connection.subscribe((event) => { delivered.push(event); }); await connection.flushBufferedSessionEvents(); expect(delivered).toEqual([{ type: "session_event", event: { type: "session_info_changed", name: "2" } }]); } finally { await connection.dispose(); } });});