@markusylisiurunen/tau 0.3.55 → 0.3.56

package/dist/host/local_session_host.js

@@ -23,15 +23,17 @@ export class LocalSessionHost {
     store;
     history;
     historyRemote;
+    onShutdown;
     clientToolBroker = new ClientToolBroker();
     sessionOptions;
     shutdownPromise;
     shuttingDown = false;
     constructor(options) {
-        const { store, history, historyRemote, ...sessionOptions } = options;
+        const { store, history, historyRemote, onShutdown, ...sessionOptions } = options;
         this.store = store;
         this.history = history;
         this.historyRemote = historyRemote;
+        this.onShutdown = onShutdown;
         this.sessionOptions = sessionOptions;
     }
     registerClientTools(options) {
@@ -99,7 +101,8 @@ export class LocalSessionHost {
         }
     }
     async createLocalSessionHandle(executionEnvironment, committedSnapshot, createParams, forceNextSnapshotRevision = false) {
-        const bootstrap = await this.resolveNewSessionBootstrap(executionEnvironment);
+        const remoteCatalog = this.sessionOptions.getRemoteModelCatalog();
+        const bootstrap = await this.resolveNewSessionBootstrap(executionEnvironment, remoteCatalog);
         if (committedSnapshot) {
             this.applySnapshotSettingsToBootstrap(bootstrap, committedSnapshot);
         }
@@ -114,9 +117,9 @@ export class LocalSessionHost {
             agentContextFiles: bootstrap.config?.agentContextFiles ?? [],
         });
         const catalog = createContentCatalogSnapshot(bootstrap);
-        return this.createLocalSessionHandleFromRuntimeContext(executionEnvironment, runtimeContext, bootstrap, catalog, committedSnapshot, createParams, forceNextSnapshotRevision);
+        return this.createLocalSessionHandleFromRuntimeContext(executionEnvironment, runtimeContext, bootstrap, catalog, remoteCatalog, committedSnapshot, createParams, forceNextSnapshotRevision);
     }
-    async createLocalSessionHandleFromRuntimeContext(executionEnvironment, runtimeContext, bootstrap, catalog, committedSnapshot, createParams, forceNextSnapshotRevision = false) {
+    async createLocalSessionHandleFromRuntimeContext(executionEnvironment, runtimeContext, bootstrap, catalog, remoteCatalog, committedSnapshot, createParams, forceNextSnapshotRevision = false) {
         let hostedSession;
         const sessionId = committedSnapshot?.sessionId ?? randomUUID();
         const createdAt = committedSnapshot?.createdAt ?? this.sessionOptions.environment.now();
@@ -130,6 +133,7 @@ export class LocalSessionHost {
             resolveSubagentPrompts: createExecutionEnvironmentSubagentPromptResolver({
                 sessionId,
                 executionEnvironment,
+                getRemoteModelCatalog: () => hostedSession.getRemoteModelCatalog(),
                 includeAgentContext: this.sessionOptions.includeAgentContext,
                 sessionStartedAt: createdAt,
             }),
@@ -153,7 +157,7 @@ export class LocalSessionHost {
         if (!attributes) {
             throw new Error("session attributes are required");
         }
-        hostedSession = new LocalHostedSessionHandle(runtime, runtimeContext.promptBootstrap, catalog, bootstrap, this.sessionOptions.includeAgentContext, executionEnvironment, this.store, this.history, this.historyRemote, attributes, createdAt, committedSnapshot, forceNextSnapshotRevision, this.sessionOptions.recordUsage ?? appendUsageLogEntry, (session) => this.sessions.delete(session));
+        hostedSession = new LocalHostedSessionHandle(runtime, runtimeContext.promptBootstrap, catalog, bootstrap, remoteCatalog, this.sessionOptions.getRemoteModelCatalog, this.sessionOptions.includeAgentContext, executionEnvironment, this.store, this.history, this.historyRemote, attributes, createdAt, committedSnapshot, forceNextSnapshotRevision, this.sessionOptions.recordUsage ?? appendUsageLogEntry, (session) => this.sessions.delete(session));
         const historyFailure = this.history.registerSession({
             sessionId: committedSnapshot?.sessionId ?? hostedSession.sessionId,
             attributes,
@@ -163,12 +167,12 @@ export class LocalSessionHost {
         this.sessions.add(hostedSession);
         return hostedSession;
     }
-    async resolveNewSessionBootstrap(executionEnvironment) {
+    async resolveNewSessionBootstrap(executionEnvironment, remoteCatalog) {
         const resolver = this.sessionOptions.resolveSessionBootstrap;
         if (!resolver) {
             return this.defaultSessionBootstrap();
         }
-        const bootstrap = await resolver({ executionEnvironment });
+        const bootstrap = await resolver({ executionEnvironment, remoteCatalog });
         return cloneResolvedBootstrap(bootstrap);
     }
     applySnapshotSettingsToBootstrap(bootstrap, snapshot) {
@@ -286,6 +290,13 @@ export class LocalSessionHost {
     }
     async shutdownNow() {
         const errors = [];
+        const resourceShutdown = this.onShutdown
+            ? Promise.resolve()
+                .then(this.onShutdown)
+                .catch((error) => {
+                errors.push(error);
+            })
+            : undefined;
         const recoveryResults = await Promise.allSettled(this.sessionRecoveryPromises.values());
         for (const result of recoveryResults) {
             if (result.status === "rejected") {
@@ -316,13 +327,19 @@ export class LocalSessionHost {
             }
         }
         this.sessions.clear();
-        await this.history.flush();
+        try {
+            await this.history.flush();
+        }
+        catch (error) {
+            errors.push(error);
+        }
         try {
             this.history.close();
         }
         catch (error) {
             errors.push(error);
         }
+        await resourceShutdown;
         if (errors.length > 0) {
             throw new AggregateError(errors, "failed to shut down local session host");
         }
@@ -354,6 +371,8 @@ class LocalHostedSessionHandle {
     promptBootstrap;
     catalog;
     bootstrap;
+    remoteCatalog;
+    getLatestRemoteModelCatalog;
     includeAgentContext;
     executionEnvironment;
     store;
@@ -397,11 +416,13 @@ class LocalHostedSessionHandle {
     pendingGoalCommit;
     forceNextSnapshotRevision;
     disposed = false;
-    constructor(runtime, promptBootstrap, catalog, bootstrap, includeAgentContext, executionEnvironment, store, history, historyRemote, attributes, createdAt, committedSnapshot, forceNextSnapshotRevision, recordUsage, removeFromHost = () => { }) {
+    constructor(runtime, promptBootstrap, catalog, bootstrap, remoteCatalog, getLatestRemoteModelCatalog, includeAgentContext, executionEnvironment, store, history, historyRemote, attributes, createdAt, committedSnapshot, forceNextSnapshotRevision, recordUsage, removeFromHost = () => { }) {
         this.runtime = runtime;
         this.promptBootstrap = promptBootstrap;
         this.catalog = catalog;
         this.bootstrap = bootstrap;
+        this.remoteCatalog = remoteCatalog;
+        this.getLatestRemoteModelCatalog = getLatestRemoteModelCatalog;
         this.includeAgentContext = includeAgentContext;
         this.executionEnvironment = executionEnvironment;
         this.store = store;
@@ -422,6 +443,9 @@ class LocalHostedSessionHandle {
         this.forceNextSnapshotRevision = forceNextSnapshotRevision;
         this.restoreProtocolState(committedSnapshot);
     }
+    getRemoteModelCatalog() {
+        return this.remoteCatalog;
+    }
     get isTurnRunning() {
         return this.runtime.isTurnRunning;
     }
@@ -870,7 +894,7 @@ class LocalHostedSessionHandle {
     }
     async setPersona(personaId) {
         this.assertActive();
-        const runtimeConfig = await this.executionEnvironment.resolveRuntimeConfig(this.executionEnvironment.snapshot().cwd);
+        const runtimeConfig = await this.executionEnvironment.resolveRuntimeConfig(this.executionEnvironment.snapshot().cwd, { remoteCatalog: this.remoteCatalog });
         const personas = runtimeConfig.personas.map((persona) => this.normalizeReloadedPersona(persona));
         const persona = personas.find((persona) => persona.id.toLowerCase() === personaId.toLowerCase());
         if (!persona) {
@@ -906,7 +930,8 @@ class LocalHostedSessionHandle {
     }
     async reload() {
         this.assertActive();
-        const runtimeConfig = await this.executionEnvironment.resolveRuntimeConfig(this.executionEnvironment.snapshot().cwd);
+        const remoteCatalog = this.getLatestRemoteModelCatalog();
+        const runtimeConfig = await this.executionEnvironment.resolveRuntimeConfig(this.executionEnvironment.snapshot().cwd, { remoteCatalog });
         if (runtimeConfig.personas.length === 0) {
             throw new Error("reload failed: no personas available");
         }
@@ -926,6 +951,7 @@ class LocalHostedSessionHandle {
             this.runtime.setPersona(nextPersona, {
                 skillsBlock: runtimeContext.promptBootstrap.promptContext.skillsBlock,
             });
+            this.remoteCatalog = remoteCatalog;
             this.bootstrap = {
                 persona: clonePersona(nextPersona),
                 discoveredSkills: structuredClone(runtimeConfig.skills),