@@ -3,6 +3,7 @@ import type { AgentContext, Connection, ConnectionContext, WSMessage } from "age
33import { isWriteCommand , safeParseEvent } from "@understudy/protocol" ;
44import type { Command , Event } from "@understudy/protocol" ;
55import { scopeSession , verifyExtensionToken } from "./auth" ;
6+ import { SESSION_NOT_CONNECTED } from "./coordinator" ;
67import { CfSessionCoordinator } from "./coordinator-cf" ;
78import { resolveSecret } from "./secrets" ;
89import type { Env , SessionState , SessionStatus } from "./types" ;
@@ -35,6 +36,7 @@ export class SessionAgent extends Agent<Env, SessionState> {
3536 if ( this . isAuthorizedConnection ( connection ) ) connection . send ( payload ) ;
3637 }
3738 } ,
39+ hasAuthorizedConnection : ( ) => this . hasAuthorizedConnection ( ) ,
3840 getAwaitingCommandIds : ( ) => this . state . awaitingCommandIds ,
3941 persistAwaitingCommandIds : ( ids ) => this . setState ( { ...this . state , awaitingCommandIds : ids } ) ,
4042 persistStatus : ( status ) => this . setState ( { ...this . state , status } ) ,
@@ -125,7 +127,14 @@ export class SessionAgent extends Agent<Env, SessionState> {
125127 }
126128
127129 async onClose ( connection : Connection , code : number , reason : string , wasClean : boolean ) : Promise < void > {
128- this . setState ( { ...this . state , status : "detached" } ) ;
130+ // Only the LAST authorized socket's close detaches the session: a late
131+ // close event from a replaced socket must not stamp "detached" over a
132+ // healthy reconnect (the closing connection is excluded explicitly, in
133+ // case the SDK has not yet reaped it from the connection set).
134+ const stillLive = [ ...this . getConnections ( ) ] . some (
135+ ( c ) => c !== connection && this . isAuthorizedConnection ( c ) ,
136+ ) ;
137+ if ( ! stillLive ) this . setState ( { ...this . state , status : "detached" } ) ;
129138 }
130139
131140 async dispatch ( command : Command , dryRun ?: boolean ) : Promise < Event > {
@@ -136,26 +145,20 @@ export class SessionAgent extends Agent<Env, SessionState> {
136145 return this . coordinator . send ( command ) ;
137146 }
138147
139- const ok = await this . checkRefResolves ( this . commandRef ( command ) ) ;
140- return {
141- type : "action_result" ,
142- commandId : command . commandId ,
143- ok,
144- ...( ok ? { } : { error : "dry-run: ref did not resolve" } ) ,
145- simulated : true ,
146- } ;
148+ const probe = await this . checkRefResolves ( this . commandRef ( command ) ) ;
149+ return this . simulatedResult ( command . commandId , probe ) ;
147150 }
148151
149152 async fillSecret ( cmd : FillSecretCommand , dryRun ?: boolean ) : Promise < Event > {
150153 if ( dryRun ) {
151- const ok = await this . checkRefResolves ( cmd . ref ) ;
152- return {
153- type : "action_result" ,
154- commandId : cmd . commandId ,
155- ok ,
156- ... ( ok ? { } : { error : "dry-run: ref did not resolve" } ) ,
157- simulated : true ,
158- } ;
154+ return this . simulatedResult ( cmd . commandId , await this . checkRefResolves ( cmd . ref ) ) ;
155+ }
156+
157+ // Gate BEFORE the vault: resolving a secret for a command that cannot
158+ // dispatch would materialize plaintext (and emit a vault access) for
159+ // nothing - fail-fast matters most exactly here (DL-004).
160+ if ( ! this . hasAuthorizedConnection ( ) ) {
161+ throw new Error ( ` ${ SESSION_NOT_CONNECTED } : no authorized extension connection` ) ;
159162 }
160163
161164 let secret : string ;
@@ -198,15 +201,37 @@ export class SessionAgent extends Agent<Env, SessionState> {
198201 // (generation bump), so it can never contain the consumer's ref AND it
199202 // invalidates the consumer's outstanding refs, breaking the approved
200203 // command that follows the dry-run.
201- private async checkRefResolves ( ref : string | undefined ) : Promise < boolean > {
202- if ( ref === undefined ) return true ;
204+ private async checkRefResolves (
205+ ref : string | undefined ,
206+ ) : Promise < { ok : true } | { ok : false ; reason : string } > {
207+ if ( ref === undefined ) return { ok : true } ;
203208
204209 const ev = await this . coordinator . send ( {
205210 type : "resolve_ref" ,
206211 commandId : crypto . randomUUID ( ) ,
207212 ref,
208213 } ) ;
209- return ev . type === "action_result" && ev . ok ;
214+ if ( ev . type !== "action_result" ) {
215+ return { ok : false , reason : `unexpected probe response '${ ev . type } '` } ;
216+ }
217+ if ( ev . ok ) return { ok : true } ;
218+ // Surface the extension's own reason (e.g. "stale or unknown ref: s1e2")
219+ // instead of collapsing every probe failure into one generic string.
220+ return { ok : false , reason : ev . error ?? "ref did not resolve" } ;
221+ }
222+
223+ /** The simulated action_result a dry-run returns in place of dispatching. */
224+ private simulatedResult (
225+ commandId : string ,
226+ probe : { ok : true } | { ok : false ; reason : string } ,
227+ ) : Event {
228+ return {
229+ type : "action_result" ,
230+ commandId,
231+ ok : probe . ok ,
232+ ...( probe . ok ? { } : { error : `dry-run: ${ probe . reason } ` } ) ,
233+ simulated : true ,
234+ } ;
210235 }
211236
212237 private commandRef ( command : Command ) : string | undefined {
@@ -224,4 +249,14 @@ export class SessionAgent extends Agent<Env, SessionState> {
224249 private isAuthorizedConnection ( connection : Connection ) : boolean {
225250 return ( connection . state as { authorized ?: boolean } | null ) ?. authorized === true ;
226251 }
252+
253+ // The delivery predicate the coordinator's fail-fast gate consults: the
254+ // same precondition sendToExtension relies on, NOT the persisted status
255+ // scalar (which a late onClose from a replaced socket can leave stale).
256+ private hasAuthorizedConnection ( ) : boolean {
257+ for ( const connection of this . getConnections ( ) ) {
258+ if ( this . isAuthorizedConnection ( connection ) ) return true ;
259+ }
260+ return false ;
261+ }
227262}
0 commit comments