@jupyter-ai/persona-manager 0.2.0-a6 → 0.2.0

package/lib/index.js

@@ -1,5 +1,4 @@
 import { IChatCommandRegistry, IInputToolbarRegistryFactory, InputToolbarRegistry } from '@jupyter/chat';
-import { IEventListener } from 'jupyterlab-eventlistener';
 import { PersonaControls } from './persona-controls';
 import { IPersonaControlRegistry, PersonaControlRegistry } from './persona-control-registry';
 import { IPersonaSessionRegistry, PersonaSessionRegistry } from './persona-events';
@@ -24,16 +23,15 @@ const plugin = {
 };
 /**
  * Plugin that provides the shared per-chat persona session-state registry,
- * fed by Jupyter Events via `jupyterlab-eventlistener`.
+ * fed by Jupyter Events via the `ServiceManager` event bus.
  */
 const sessionRegistryPlugin = {
     id: '@jupyter-ai/persona-manager:session-registry',
     description: 'Provides the per-chat persona session-state registry (fed by Jupyter Events).',
     autoStart: true,
     provides: IPersonaSessionRegistry,
-    requires: [IEventListener],
-    activate: (app, eventListener) => {
-        return new PersonaSessionRegistry(eventListener);
+    activate: (app) => {
+        return new PersonaSessionRegistry(app.serviceManager.events);
     }
 };
 /**
@@ -87,6 +85,9 @@ const toolbarPlugin = {
             sessionRegistry,
             events
         });
+        // Wrap the stop button to inject the session registry, so it can enable
+        // itself while a persona in the chat is processing a message.
+        const StopButtonItem = (itemProps) => StopButton({ ...itemProps, sessionRegistry });
         return {
             create: () => {
                 // Start with the default toolbar (Send, Attach, Cancel, SaveEdit)
@@ -100,7 +101,7 @@ const toolbarPlugin = {
                 // toolbar's cancel button (10); a tie would leave the order to
                 // insertion rather than position.
                 registry.addItem('stop', {
-                    element: StopButton,
+                    element: StopButtonItem,
                     position: 7
                 });
                 return registry;

package/lib/persona-events.d.ts

@@ -1,4 +1,17 @@
-import { IEventListener } from 'jupyterlab-eventlistener';
+/**
+ * Frontend persona session state, delivered over Jupyter Events (via the
+ * `ServiceManager` event bus), replacing the previous Yjs-awareness channel.
+ *
+ * Flow: the server emits `personas` / `persona_state` events (each carrying the
+ * chat's stable `chat_id`); the `PersonaSessionRegistry` routes them to the
+ * per-chat `PersonaManagerSessionState`, which holds the persona list and a
+ * `PersonaSessionState` per persona and fires a Lumino `changed` signal. React
+ * components listen to that signal. When a chat closes, its state is discarded.
+ *
+ * Names mirror the Python package (`PersonaManagerSessionState`,
+ * `PersonaSessionState`).
+ */
+import { Event } from '@jupyterlab/services';
 import { Token } from '@lumino/coreutils';
 import { IDisposable } from '@lumino/disposable';
 import { ISignal } from '@lumino/signaling';
@@ -14,6 +27,7 @@ type PersonaStatePayload = {
     settings?: SettingConfiguration[];
     usage?: Usage;
     slash_commands?: CommandOption[];
+    processing?: boolean;
 };
 /**
  * One persona's live session state, built from a `persona_state` event.
@@ -27,6 +41,15 @@ export declare class PersonaSessionState {
     readonly settings: SettingConfiguration[];
     readonly usage: Usage;
     readonly slash_commands: CommandOption[];
+    /** Whether the persona is currently processing a message. */
+    readonly processing: boolean;
+    /**
+     * Merge a `persona_state` event payload onto this state, returning a new
+     * instance. Attributes the event omits are carried forward unchanged, so a
+     * partial event (e.g. usage-only) replaces only what it carries. A new
+     * reference is returned so React consumers re-render.
+     */
+    withUpdate(payload: PersonaStatePayload): PersonaSessionState;
 }
 /**
  * The per-chat manager session state: the persona list plus a
@@ -44,9 +67,18 @@ export declare class PersonaManagerSessionState implements IDisposable {
     get personas(): PersonaOption[];
     /** A persona's session state, or undefined if it has not published yet. */
     getPersona(id: string): PersonaSessionState | undefined;
+    /**
+     * Whether any persona in this chat is currently processing a message. The
+     * stop button uses this to enable itself, since a persona can be processing
+     * (thinking, running tools, awaiting an agent turn) without actively writing.
+     */
+    get processing(): boolean;
     /** Apply a `personas` event payload. */
     updatePersonas(personas: PersonaOption[]): void;
-    /** Apply a `persona_state` event payload for one persona. */
+    /**
+     * Apply a `persona_state` event payload for one persona, merging it onto the
+     * persona's existing state (absent attributes are left unchanged).
+     */
     updatePersonaState(personaId: string, payload: PersonaStatePayload): void;
     get isDisposed(): boolean;
     dispose(): void;
@@ -64,7 +96,7 @@ export declare class PersonaManagerSessionState implements IDisposable {
  * controls and the slash-command provider.
  */
 export declare class PersonaSessionRegistry {
-    constructor(eventListener: IEventListener);
+    constructor(events: Event.IManager);
     /**
      * Get (or create) the manager session state for a chat. Components call this
      * with their chat's stable id (`IChatModel.id` / `IChatContext.id`).

package/lib/persona-events.js

@@ -11,12 +11,29 @@ export const PERSONA_SELECTED_EVENT_SCHEMA_ID = 'https://schema.jupyter.org/jupy
  */
 export class PersonaSessionState {
     constructor(id, payload = {}) {
-        var _a, _b, _c, _d;
+        var _a, _b, _c, _d, _e;
         this.id = id;
         this.model = (_a = payload.model) !== null && _a !== void 0 ? _a : EMPTY_MODEL_CONFIGURATION;
         this.settings = (_b = payload.settings) !== null && _b !== void 0 ? _b : [];
         this.usage = { ...EMPTY_USAGE, ...((_c = payload.usage) !== null && _c !== void 0 ? _c : {}) };
         this.slash_commands = (_d = payload.slash_commands) !== null && _d !== void 0 ? _d : [];
+        this.processing = (_e = payload.processing) !== null && _e !== void 0 ? _e : false;
+    }
+    /**
+     * Merge a `persona_state` event payload onto this state, returning a new
+     * instance. Attributes the event omits are carried forward unchanged, so a
+     * partial event (e.g. usage-only) replaces only what it carries. A new
+     * reference is returned so React consumers re-render.
+     */
+    withUpdate(payload) {
+        var _a, _b, _c, _d, _e;
+        return new PersonaSessionState(this.id, {
+            model: (_a = payload.model) !== null && _a !== void 0 ? _a : this.model,
+            settings: (_b = payload.settings) !== null && _b !== void 0 ? _b : this.settings,
+            usage: (_c = payload.usage) !== null && _c !== void 0 ? _c : this.usage,
+            slash_commands: (_d = payload.slash_commands) !== null && _d !== void 0 ? _d : this.slash_commands,
+            processing: (_e = payload.processing) !== null && _e !== void 0 ? _e : this.processing
+        });
     }
 }
 /**
@@ -49,15 +66,35 @@ export class PersonaManagerSessionState {
     getPersona(id) {
         return this._states.get(id);
     }
+    /**
+     * Whether any persona in this chat is currently processing a message. The
+     * stop button uses this to enable itself, since a persona can be processing
+     * (thinking, running tools, awaiting an agent turn) without actively writing.
+     */
+    get processing() {
+        for (const state of this._states.values()) {
+            if (state.processing) {
+                return true;
+            }
+        }
+        return false;
+    }
     /** Apply a `personas` event payload. */
     updatePersonas(personas) {
         this._personasReceived = true;
         this._personas = personas;
         this._changed.emit();
     }
-    /** Apply a `persona_state` event payload for one persona. */
+    /**
+     * Apply a `persona_state` event payload for one persona, merging it onto the
+     * persona's existing state (absent attributes are left unchanged).
+     */
     updatePersonaState(personaId, payload) {
-        this._states.set(personaId, new PersonaSessionState(personaId, payload));
+        const prev = this._states.get(personaId);
+        const next = prev
+            ? prev.withUpdate(payload)
+            : new PersonaSessionState(personaId, payload);
+        this._states.set(personaId, next);
         this._changed.emit();
     }
     get isDisposed() {
@@ -80,15 +117,15 @@ export class PersonaManagerSessionState {
  * controls and the slash-command provider.
  */
 export class PersonaSessionRegistry {
-    constructor(eventListener) {
-        this._onPersonasEvent = async (_manager, _schemaId, event) => {
+    constructor(events) {
+        this._onPersonasEvent = async (event) => {
             const data = event;
             if (!data.chat_id) {
                 return;
             }
             this.get(data.chat_id).updatePersonas(Array.isArray(data.personas) ? data.personas : []);
         };
-        this._onPersonaStateEvent = async (_manager, _schemaId, event) => {
+        this._onPersonaStateEvent = async (event) => {
             const data = event;
             if (!data.chat_id || !data.persona_id) {
                 return;
@@ -96,8 +133,18 @@ export class PersonaSessionRegistry {
             this.get(data.chat_id).updatePersonaState(data.persona_id, data);
         };
         this._byChatId = new Map();
-        eventListener.addListener(PERSONAS_EVENT_SCHEMA_ID, this._onPersonasEvent);
-        eventListener.addListener(PERSONA_STATE_EVENT_SCHEMA_ID, this._onPersonaStateEvent);
+        // The ServiceManager event bus (JupyterLab >= 4.0) exposes a single shared
+        // stream of all Jupyter Events; filter it by schema id to route the two
+        // persona event types. This supersedes the former `jupyterlab-eventlistener`
+        // dependency.
+        events.stream.connect((_, emission) => {
+            if (emission.schema_id === PERSONAS_EVENT_SCHEMA_ID) {
+                void this._onPersonasEvent(emission);
+            }
+            else if (emission.schema_id === PERSONA_STATE_EVENT_SCHEMA_ID) {
+                void this._onPersonaStateEvent(emission);
+            }
+        });
     }
     /**
      * Get (or create) the manager session state for a chat. Components call this

package/lib/stop-button.d.ts

@@ -1,7 +1,19 @@
 import { InputToolbarRegistry } from '@jupyter/chat';
+import { PersonaSessionRegistry } from './persona-events';
 /**
- * A stop button for the chat input toolbar. Observes the chat model's
- * writers list to enable itself when an AI persona is actively writing,
- * and calls the persona cancel endpoint on click.
+ * A stop button for the chat input toolbar. Enables itself while a persona in
+ * the chat is processing a message — tracked via the persona-state events fed
+ * into the `PersonaSessionRegistry` — and calls the persona cancel endpoint on
+ * click.
+ *
+ * This tracks *processing* rather than the chat's writers list: a persona can
+ * be processing (thinking, running tools, awaiting an agent turn) without
+ * actively writing, and should be interruptible the whole time.
  */
-export declare function StopButton(props: InputToolbarRegistry.IToolbarItemProps): JSX.Element;
+export declare function StopButton(props: InputToolbarRegistry.IToolbarItemProps & {
+    /**
+     * The per-chat persona session-state registry (fed by Jupyter Events).
+     * Optional so the button still renders (permanently disabled) without it.
+     */
+    sessionRegistry?: PersonaSessionRegistry;
+}): JSX.Element;

package/lib/stop-button.js

@@ -4,39 +4,62 @@ import { TooltippedIconButton } from '@jupyter/chat';
 import { cancelResponse } from './request';
 const STOP_BUTTON_CLASS = 'jp-jai-stopButton';
 /**
- * A stop button for the chat input toolbar. Observes the chat model's
- * writers list to enable itself when an AI persona is actively writing,
- * and calls the persona cancel endpoint on click.
+ * A stop button for the chat input toolbar. Enables itself while a persona in
+ * the chat is processing a message — tracked via the persona-state events fed
+ * into the `PersonaSessionRegistry` — and calls the persona cancel endpoint on
+ * click.
+ *
+ * This tracks *processing* rather than the chat's writers list: a persona can
+ * be processing (thinking, running tools, awaiting an agent turn) without
+ * actively writing, and should be interruptible the whole time.
  */
 export function StopButton(props) {
-    const { chatModel } = props;
+    var _a;
+    const { chatModel, sessionRegistry } = props;
     const [disabled, setDisabled] = useState(true);
     const [inFlight, setInFlight] = useState(false);
     const tooltip = 'Stop generating';
+    // The chat's stable id is assigned asynchronously (once the model is
+    // `ready`), so track it in state and update it when the model becomes ready.
+    const [chatId, setChatId] = useState((_a = chatModel === null || chatModel === void 0 ? void 0 : chatModel.id) !== null && _a !== void 0 ? _a : null);
     useEffect(() => {
-        var _a;
         if (!chatModel) {
-            setDisabled(true);
+            setChatId(null);
             return;
         }
-        const checkWriters = () => {
-            const hasAIWriter = chatModel.writers.some(w => w.user.bot);
-            setDisabled(!hasAIWriter);
-        };
-        checkWriters();
-        (_a = chatModel.writersChanged) === null || _a === void 0 ? void 0 : _a.connect(checkWriters);
+        let cancelled = false;
+        void chatModel.ready.then(id => {
+            if (!cancelled) {
+                setChatId(id !== null && id !== void 0 ? id : null);
+            }
+        });
         return () => {
-            var _a;
-            (_a = chatModel.writersChanged) === null || _a === void 0 ? void 0 : _a.disconnect(checkWriters);
+            cancelled = true;
         };
     }, [chatModel]);
+    // Enable the button while any persona in this chat is processing. The
+    // manager session state fires `changed` on every persona-state update, which
+    // includes the processing transitions emitted by `track_processing`.
+    useEffect(() => {
+        if (!sessionRegistry || !chatId) {
+            setDisabled(true);
+            return;
+        }
+        const managerState = sessionRegistry.get(chatId);
+        const sync = () => setDisabled(!managerState.processing);
+        sync();
+        managerState.changed.connect(sync);
+        return () => {
+            managerState.changed.disconnect(sync);
+        };
+    }, [sessionRegistry, chatId]);
     async function handleStop() {
-        if (!(chatModel === null || chatModel === void 0 ? void 0 : chatModel.id)) {
+        if (!chatId) {
             return;
         }
         setInFlight(true);
         try {
-            await cancelResponse(chatModel.id);
+            await cancelResponse(chatId);
         }
         finally {
             setInFlight(false);

package/package.json

@@ -1,6 +1,6 @@
 {
     "name": "@jupyter-ai/persona-manager",
-    "version": "0.2.0-a6",
+    "version": "0.2.0",
     "description": "The core manager & registry for AI personas in Jupyter AI",
     "keywords": [
         "jupyter",
@@ -61,8 +61,7 @@
         "@jupyter/chat": "^0.25.0-alpha.5",
         "@jupyterlab/application": "^4.0.0",
         "@jupyterlab/coreutils": "^6.0.0",
-        "@jupyterlab/services": "^7.0.0",
-        "jupyterlab-eventlistener": "^0.4.0"
+        "@jupyterlab/services": "^7.0.0"
     },
     "devDependencies": {
         "@jupyterlab/builder": "^4.0.0",

package/src/__tests__/persona-events.spec.ts

@@ -3,7 +3,8 @@
  * per-chat routing, the `changed` signal, and discarding a chat's state when it
  * closes.
  */
-import { IEventListener } from 'jupyterlab-eventlistener';
+import { Event } from '@jupyterlab/services';
+import { Signal } from '@lumino/signaling';
 
 import {
   PERSONAS_EVENT_SCHEMA_ID,
@@ -11,31 +12,26 @@ import {
   PersonaSessionRegistry
 } from '../persona-events';
 
-type Listener = (m: unknown, s: string, e: any) => Promise<void>;
-
-/** A stand-in for IEventListener that lets tests emit events to listeners. */
-class FakeEventListener {
-  private _bySchema = new Map<string, Listener[]>();
+/**
+ * A stand-in for the ServiceManager event bus (`Event.IManager`) that lets
+ * tests push events onto the shared stream the registry subscribes to.
+ */
+class FakeEventManager {
+  readonly stream = new Signal<this, Event.Emission>(this);
 
-  addListener(schemaId: string, listener: Listener): void {
-    const list = this._bySchema.get(schemaId) ?? [];
-    list.push(listener);
-    this._bySchema.set(schemaId, list);
-  }
+  /** Push an event onto the stream, mirroring the jupyter_server event bus. */
   async emit(schemaId: string, data: any): Promise<void> {
-    for (const l of this._bySchema.get(schemaId) ?? []) {
-      await l(null, schemaId, data);
-    }
+    this.stream.emit({ schema_id: schemaId, ...data });
   }
 }
 
 function makeRegistry(): {
   registry: PersonaSessionRegistry;
-  events: FakeEventListener;
+  events: FakeEventManager;
 } {
-  const events = new FakeEventListener();
+  const events = new FakeEventManager();
   const registry = new PersonaSessionRegistry(
-    events as unknown as IEventListener
+    events as unknown as Event.IManager
   );
   return { registry, events };
 }
@@ -66,6 +62,38 @@ describe('PersonaSessionRegistry', () => {
     expect(state?.usage.input_tokens).toBe(5);
   });
 
+  it('merges partial events, leaving unset attributes unchanged', async () => {
+    const { registry, events } = makeRegistry();
+    // First event sets the model.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1',
+      model: { current: 'm1', options: [], settings: [] }
+    });
+    // A usage-only event must not drop the previously-set model.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1',
+      usage: { input_tokens: 5 }
+    });
+    const state = registry.get('a.chat').getPersona('p1');
+    expect(state?.model.current).toBe('m1');
+    expect(state?.usage.input_tokens).toBe(5);
+
+    // A processing-only event likewise preserves model and usage.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1',
+      processing: true
+    });
+    const after = registry.get('a.chat').getPersona('p1');
+    expect(after?.model.current).toBe('m1');
+    expect(after?.usage.input_tokens).toBe(5);
+    expect(after?.processing).toBe(true);
+    // A new instance is produced on each merge, so React consumers re-render.
+    expect(after).not.toBe(state);
+  });
+
   it('fires the changed signal on updates', async () => {
     const { registry, events } = makeRegistry();
     const managerState = registry.get('a.chat');
@@ -84,6 +112,47 @@ describe('PersonaSessionRegistry', () => {
     expect(fired).toBe(2);
   });
 
+  it('tracks per-persona processing and exposes chat-level processing', async () => {
+    const { registry, events } = makeRegistry();
+    const managerState = registry.get('a.chat');
+    expect(managerState.processing).toBe(false);
+
+    // A persona reports it started processing.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1',
+      processing: true
+    });
+    expect(managerState.getPersona('p1')?.processing).toBe(true);
+    expect(managerState.processing).toBe(true);
+
+    // A second, idle persona doesn't flip the chat back to idle.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p2',
+      processing: false
+    });
+    expect(managerState.processing).toBe(true);
+
+    // The processing persona reports it finished; chat is now idle.
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1',
+      processing: false
+    });
+    expect(managerState.processing).toBe(false);
+  });
+
+  it('defaults persona processing to false when the field is absent', async () => {
+    const { registry, events } = makeRegistry();
+    await events.emit(PERSONA_STATE_EVENT_SCHEMA_ID, {
+      chat_id: 'a.chat',
+      persona_id: 'p1'
+    });
+    expect(registry.get('a.chat').getPersona('p1')?.processing).toBe(false);
+    expect(registry.get('a.chat').processing).toBe(false);
+  });
+
   it('discards a chat session state on close, freeing memory', async () => {
     const { registry, events } = makeRegistry();
     await events.emit(PERSONAS_EVENT_SCHEMA_ID, {

package/src/index.ts

@@ -9,8 +9,6 @@ import {
   InputToolbarRegistry
 } from '@jupyter/chat';
 
-import { IEventListener } from 'jupyterlab-eventlistener';
-
 import { PersonaControls } from './persona-controls';
 import {
   IPersonaControlRegistry,
@@ -51,7 +49,7 @@ const plugin: JupyterFrontEndPlugin<void> = {
 
 /**
  * Plugin that provides the shared per-chat persona session-state registry,
- * fed by Jupyter Events via `jupyterlab-eventlistener`.
+ * fed by Jupyter Events via the `ServiceManager` event bus.
  */
 const sessionRegistryPlugin: JupyterFrontEndPlugin<PersonaSessionRegistry> = {
   id: '@jupyter-ai/persona-manager:session-registry',
@@ -59,12 +57,8 @@ const sessionRegistryPlugin: JupyterFrontEndPlugin<PersonaSessionRegistry> = {
     'Provides the per-chat persona session-state registry (fed by Jupyter Events).',
   autoStart: true,
   provides: IPersonaSessionRegistry,
-  requires: [IEventListener],
-  activate: (
-    app: JupyterFrontEnd,
-    eventListener: IEventListener
-  ): PersonaSessionRegistry => {
-    return new PersonaSessionRegistry(eventListener);
+  activate: (app: JupyterFrontEnd): PersonaSessionRegistry => {
+    return new PersonaSessionRegistry(app.serviceManager.events);
   }
 };
 
@@ -133,6 +127,11 @@ const toolbarPlugin: JupyterFrontEndPlugin<IInputToolbarRegistryFactory> = {
         sessionRegistry,
         events
       });
+    // Wrap the stop button to inject the session registry, so it can enable
+    // itself while a persona in the chat is processing a message.
+    const StopButtonItem = (
+      itemProps: InputToolbarRegistry.IToolbarItemProps
+    ) => StopButton({ ...itemProps, sessionRegistry });
     return {
       create: () => {
         // Start with the default toolbar (Send, Attach, Cancel, SaveEdit)
@@ -146,7 +145,7 @@ const toolbarPlugin: JupyterFrontEndPlugin<IInputToolbarRegistryFactory> = {
         // toolbar's cancel button (10); a tie would leave the order to
         // insertion rather than position.
         registry.addItem('stop', {
-          element: StopButton,
+          element: StopButtonItem,
           position: 7
         });
         return registry;

package/src/persona-events.ts

@@ -1,6 +1,6 @@
 /**
- * Frontend persona session state, delivered over Jupyter Events (via
- * `jupyterlab-eventlistener`), replacing the previous Yjs-awareness channel.
+ * Frontend persona session state, delivered over Jupyter Events (via the
+ * `ServiceManager` event bus), replacing the previous Yjs-awareness channel.
  *
  * Flow: the server emits `personas` / `persona_state` events (each carrying the
  * chat's stable `chat_id`); the `PersonaSessionRegistry` routes them to the
@@ -12,7 +12,6 @@
  * `PersonaSessionState`).
  */
 import { Event } from '@jupyterlab/services';
-import { IEventListener } from 'jupyterlab-eventlistener';
 import { Token } from '@lumino/coreutils';
 import { IDisposable } from '@lumino/disposable';
 import { ISignal, Signal } from '@lumino/signaling';
@@ -42,6 +41,7 @@ type PersonaStatePayload = {
   settings?: SettingConfiguration[];
   usage?: Usage;
   slash_commands?: CommandOption[];
+  processing?: boolean;
 };
 
 /** The wire shape of a `personas` event. */
@@ -64,12 +64,31 @@ export class PersonaSessionState {
     this.settings = payload.settings ?? [];
     this.usage = { ...EMPTY_USAGE, ...(payload.usage ?? {}) };
     this.slash_commands = payload.slash_commands ?? [];
+    this.processing = payload.processing ?? false;
   }
 
   readonly model: ModelConfiguration;
   readonly settings: SettingConfiguration[];
   readonly usage: Usage;
   readonly slash_commands: CommandOption[];
+  /** Whether the persona is currently processing a message. */
+  readonly processing: boolean;
+
+  /**
+   * Merge a `persona_state` event payload onto this state, returning a new
+   * instance. Attributes the event omits are carried forward unchanged, so a
+   * partial event (e.g. usage-only) replaces only what it carries. A new
+   * reference is returned so React consumers re-render.
+   */
+  withUpdate(payload: PersonaStatePayload): PersonaSessionState {
+    return new PersonaSessionState(this.id, {
+      model: payload.model ?? this.model,
+      settings: payload.settings ?? this.settings,
+      usage: payload.usage ?? this.usage,
+      slash_commands: payload.slash_commands ?? this.slash_commands,
+      processing: payload.processing ?? this.processing
+    });
+  }
 }
 
 /**
@@ -100,6 +119,20 @@ export class PersonaManagerSessionState implements IDisposable {
     return this._states.get(id);
   }
 
+  /**
+   * Whether any persona in this chat is currently processing a message. The
+   * stop button uses this to enable itself, since a persona can be processing
+   * (thinking, running tools, awaiting an agent turn) without actively writing.
+   */
+  get processing(): boolean {
+    for (const state of this._states.values()) {
+      if (state.processing) {
+        return true;
+      }
+    }
+    return false;
+  }
+
   /** Apply a `personas` event payload. */
   updatePersonas(personas: PersonaOption[]): void {
     this._personasReceived = true;
@@ -107,9 +140,16 @@ export class PersonaManagerSessionState implements IDisposable {
     this._changed.emit();
   }
 
-  /** Apply a `persona_state` event payload for one persona. */
+  /**
+   * Apply a `persona_state` event payload for one persona, merging it onto the
+   * persona's existing state (absent attributes are left unchanged).
+   */
   updatePersonaState(personaId: string, payload: PersonaStatePayload): void {
-    this._states.set(personaId, new PersonaSessionState(personaId, payload));
+    const prev = this._states.get(personaId);
+    const next = prev
+      ? prev.withUpdate(payload)
+      : new PersonaSessionState(personaId, payload);
+    this._states.set(personaId, next);
     this._changed.emit();
   }
 
@@ -141,15 +181,18 @@ export class PersonaManagerSessionState implements IDisposable {
  * controls and the slash-command provider.
  */
 export class PersonaSessionRegistry {
-  constructor(eventListener: IEventListener) {
-    eventListener.addListener(
-      PERSONAS_EVENT_SCHEMA_ID,
-      this._onPersonasEvent as any
-    );
-    eventListener.addListener(
-      PERSONA_STATE_EVENT_SCHEMA_ID,
-      this._onPersonaStateEvent as any
-    );
+  constructor(events: Event.IManager) {
+    // The ServiceManager event bus (JupyterLab >= 4.0) exposes a single shared
+    // stream of all Jupyter Events; filter it by schema id to route the two
+    // persona event types. This supersedes the former `jupyterlab-eventlistener`
+    // dependency.
+    events.stream.connect((_, emission) => {
+      if (emission.schema_id === PERSONAS_EVENT_SCHEMA_ID) {
+        void this._onPersonasEvent(emission);
+      } else if (emission.schema_id === PERSONA_STATE_EVENT_SCHEMA_ID) {
+        void this._onPersonaStateEvent(emission);
+      }
+    });
   }
 
   /**
@@ -182,11 +225,7 @@ export class PersonaSessionRegistry {
     }
   }
 
-  private _onPersonasEvent = async (
-    _manager: unknown,
-    _schemaId: string,
-    event: Event.Emission
-  ): Promise<void> => {
+  private _onPersonasEvent = async (event: Event.Emission): Promise<void> => {
     const data = event as PersonasPayload;
     if (!data.chat_id) {
       return;
@@ -197,8 +236,6 @@ export class PersonaSessionRegistry {
   };
 
   private _onPersonaStateEvent = async (
-    _manager: unknown,
-    _schemaId: string,
     event: Event.Emission
   ): Promise<void> => {
     const data = event as PersonaStatePayload;

package/src/stop-button.tsx

@@ -2,49 +2,78 @@ import React, { useEffect, useState } from 'react';
 import StopIcon from '@mui/icons-material/Stop';
 import { InputToolbarRegistry, TooltippedIconButton } from '@jupyter/chat';
 import { cancelResponse } from './request';
+import { PersonaSessionRegistry } from './persona-events';
 
 const STOP_BUTTON_CLASS = 'jp-jai-stopButton';
 
 /**
- * A stop button for the chat input toolbar. Observes the chat model's
- * writers list to enable itself when an AI persona is actively writing,
- * and calls the persona cancel endpoint on click.
+ * A stop button for the chat input toolbar. Enables itself while a persona in
+ * the chat is processing a message — tracked via the persona-state events fed
+ * into the `PersonaSessionRegistry` — and calls the persona cancel endpoint on
+ * click.
+ *
+ * This tracks *processing* rather than the chat's writers list: a persona can
+ * be processing (thinking, running tools, awaiting an agent turn) without
+ * actively writing, and should be interruptible the whole time.
  */
 export function StopButton(
-  props: InputToolbarRegistry.IToolbarItemProps
+  props: InputToolbarRegistry.IToolbarItemProps & {
+    /**
+     * The per-chat persona session-state registry (fed by Jupyter Events).
+     * Optional so the button still renders (permanently disabled) without it.
+     */
+    sessionRegistry?: PersonaSessionRegistry;
+  }
 ): JSX.Element {
-  const { chatModel } = props;
+  const { chatModel, sessionRegistry } = props;
   const [disabled, setDisabled] = useState(true);
   const [inFlight, setInFlight] = useState(false);
   const tooltip = 'Stop generating';
 
+  // The chat's stable id is assigned asynchronously (once the model is
+  // `ready`), so track it in state and update it when the model becomes ready.
+  const [chatId, setChatId] = useState<string | null>(chatModel?.id ?? null);
   useEffect(() => {
     if (!chatModel) {
-      setDisabled(true);
+      setChatId(null);
       return;
     }
-
-    const checkWriters = () => {
-      const hasAIWriter = chatModel.writers.some(w => w.user.bot);
-      setDisabled(!hasAIWriter);
+    let cancelled = false;
+    void chatModel.ready.then(id => {
+      if (!cancelled) {
+        setChatId(id ?? null);
+      }
+    });
+    return () => {
+      cancelled = true;
     };
+  }, [chatModel]);
 
-    checkWriters();
-    chatModel.writersChanged?.connect(checkWriters);
-
+  // Enable the button while any persona in this chat is processing. The
+  // manager session state fires `changed` on every persona-state update, which
+  // includes the processing transitions emitted by `track_processing`.
+  useEffect(() => {
+    if (!sessionRegistry || !chatId) {
+      setDisabled(true);
+      return;
+    }
+    const managerState = sessionRegistry.get(chatId);
+    const sync = () => setDisabled(!managerState.processing);
+    sync();
+    managerState.changed.connect(sync);
     return () => {
-      chatModel.writersChanged?.disconnect(checkWriters);
+      managerState.changed.disconnect(sync);
     };
-  }, [chatModel]);
+  }, [sessionRegistry, chatId]);
 
   async function handleStop() {
-    if (!chatModel?.id) {
+    if (!chatId) {
       return;
     }
 
     setInFlight(true);
     try {
-      await cancelResponse(chatModel.id);
+      await cancelResponse(chatId);
     } finally {
       setInFlight(false);
     }