Sessions: re-apply tool restriction after session/end-seed to recover plugin tools

src/app/sessions.ts

@@ -101,6 +101,8 @@ export class Sessions {
   private readonly bindings = new Map<Agent, SessionBinding>()
   private readonly retryHealth = new RetryHealth()
   private readonly ctx: Context
+  /** Restrict disposers that must be re-applied after a `session/end-seed` resets the agent's scope. */
+  private readonly restrictDisposers = new Map<Agent, () => void>()
   /** Registers ordered `relay` and `notice` effect handlers that deliver to the main Session. */
   constructor(private readonly host: SessionsHost) {
     this.ctx = host.context
@@ -122,6 +124,10 @@ export class Sessions {
       if (!agent || !binding || binding.subagent) return
       this.retryHealth.observe(agent.id, event)
       this.updateModelHealth()
+      // After a process restart DSH writes `session/end-seed` and rebuilds the agent's scope. Plugin tools may not yet be
+      // re-registered at that point, so the restrict filter that was applied during `decorate` can silently drop them.
+      // Re-apply the restriction once per seed so the next request picks up any tools that arrived in the meantime.
+      if ((event as { type: string }).type === 'session/end-seed') this.reapplyRestrict(agent)
     })
     ctx.on('agent/created', async ({ agent }) => { await this.decorate(agent); return undefined })
     // Installing happens after setup, which takes a moment, so a Session can reach DSH before this subscription exists: see `restore`.
@@ -295,9 +301,8 @@ export class Sessions {
     const register = this.registerTool(agent, binding)
     for (const tool of host.tools(agent, ticket)) register(tool.name, tool.description, tool.schema, tool.execute)
     if (master) {
-      binding.disposers.push(agent.ctx.tools.restrict({ allow: ['skill', SHELL_TOOL, 'read', 'read_image', 'job_output', 'job_list', 'job_kill',
-        'workbuddy_web_search', ...CONTEXT_TOOLS].filter(name => agent.ctx.tools.get(name, scopeParentOf(agent))) }),
-        agent.ctx.tools.guard(execution => noWait(execution, () => this.routed(agent))))
+      binding.disposers.push(agent.ctx.tools.guard(execution => noWait(execution, () => this.routed(agent))))
+      this.applyRestrict(agent, binding)
     } else {
       // An Executor Session keeps the profile's work tools it has, which DSH requires to be known; its own tools are not masked.
       binding.disposers.push(...this.workOnly(agent, ticket!.id))
@@ -331,6 +336,28 @@ export class Sessions {
       return !current || !liveTicket(current) ? 'Ticket ended; new work needs a new authorized ticket' : undefined
     })]
   }
+  /**
+   * Applies (or re-applies) the master Session's tool restriction. Plugin tools like billion-context's may not be registered
+   * when `decorate` first runs or after a `session/end-seed` resets the scope, so the filter is recomputed each time and
+   * unknown names are silently dropped rather than crashing `restrict`. The previous restrict disposer is lifted first.
+   */
+  private applyRestrict(agent: Agent, binding: SessionBinding): void {
+    this.restrictDisposers.get(agent)?.()
+    const allow = ['skill', SHELL_TOOL, 'read', 'read_image', 'job_output', 'job_list', 'job_kill',
+      'workbuddy_web_search', ...CONTEXT_TOOLS].filter(name => agent.ctx.tools.get(name, scopeParentOf(agent)))
+    try {
+      this.restrictDisposers.set(agent, agent.ctx.tools.restrict({ allow }))
+      binding.disposers.push(() => { this.restrictDisposers.delete(agent) })
+    } catch (error) {
+      this.ctx.logger.warn('dsh-superagent could not apply tool restriction to %s: %s', agent.id, String(error))
+    }
+  }
+  /** Re-applies the master Session's tool restriction after a `session/end-seed` rebuilt its scope. */
+  private reapplyRestrict(agent: Agent): void {
+    const binding = this.bindings.get(agent)
+    if (!binding || binding.subagent) return
+    this.applyRestrict(agent, binding)
+  }
   /**
    * A subagent of an Executor Session, or of one of its subagents, which runs under DSH's `standard` preset with the worker's
    * permissions, gets the same work tools as the worker, and none once the worker's ticket ended. Any other agent is left as