Fix always-on skill recovery

Generated-By: PostHog Code
Task-Id: 73390913-7191-4450-9fd3-31dd12879508

Author: adboio

packages/agent/src/server/agent-server.test.ts

@@ -2204,14 +2204,13 @@ describe("AgentServer HTTP Mode", () => {
       )?.text;
 
       expect(sentText).toBe("/local-test-skill with context");
-      expect(sentMeta?.localSkillContext).toContain(
-        'local skill "/local-test-skill"',
-      );
       expect(sentMeta?.localSkillContext).toContain("LOCAL_SKILL_MARKER");
       expect(sentMeta?.localSkillContext).toContain(
         "always-on skills apply for the entire session",
       );
-      expect(sentMeta?.localSkillContext).toContain("with context");
+      expect(
+        String(sentMeta?.localSkillContext).match(/LOCAL_SKILL_MARKER/g),
+      ).toHaveLength(1);
       expect(sentMeta?.localSkillName).toBe("local-test-skill");
     }, 20000);
 

packages/agent/src/server/agent-server.ts

@@ -2734,6 +2734,14 @@ export class AgentServer {
         : null;
 
     if (invocation) {
+      if (
+        alwaysOnSkills.some((skill) => skill.skillName === invocation.skillName)
+      ) {
+        return {
+          skillName: invocation.skillName,
+          context: alwaysOnContext ?? "",
+        };
+      }
       const hasMatchingArtifact = artifacts.some(
         (artifact) =>
           artifact.type === "skill_bundle" &&
@@ -2768,7 +2776,7 @@ export class AgentServer {
       .join("\n");
     const attachedContext = this.buildAttachedSkillsPromptContext(
       runId,
-      artifacts,
+      artifacts.filter((artifact) => artifact.metadata?.always_on !== true),
       messageText,
     );
     if (!alwaysOnContext) return attachedContext;

packages/core/src/sessions/sessionService.ts

@@ -348,6 +348,14 @@ export interface SessionServiceDeps {
     options: SessionConfigOption[],
   ) => void;
   removePersistedConfigOptions: (taskRunId: string) => void;
+  getPersistedAlwaysOnSkillInstructions?: (
+    taskRunId: string,
+  ) => string | undefined;
+  setPersistedAlwaysOnSkillInstructions?: (
+    taskRunId: string,
+    instructions: string,
+  ) => void;
+  removePersistedAlwaysOnSkillInstructions?: (taskRunId: string) => void;
   adapterStore: {
     getAdapter(taskRunId: string): Adapter | undefined;
     setAdapter(taskRunId: string, adapter: Adapter): void;
@@ -1949,6 +1957,10 @@ export class SessionService {
     const persistedConfigOptions = this.d.getPersistedConfigOptions(taskRunId);
 
     const previous = this.d.store.getSessions()[taskRunId];
+    const resolvedAlwaysOnSkillInstructions =
+      alwaysOnSkillInstructions ??
+      previous?.alwaysOnSkillInstructions ??
+      this.d.getPersistedAlwaysOnSkillInstructions?.(taskRunId);
 
     const session = createBaseSession(taskRunId, taskId, taskTitle);
     // Repainting from the log must not blank a transcript we already hold:
@@ -2019,7 +2031,7 @@ export class SessionService {
         this.d.settings;
       const effectiveCustomInstructions = [
         customInstructions,
-        alwaysOnSkillInstructions ?? previous?.alwaysOnSkillInstructions,
+        resolvedAlwaysOnSkillInstructions,
       ]
         .filter((value): value is string => Boolean(value))
         .join("\n\n");
@@ -2040,8 +2052,7 @@ export class SessionService {
       });
 
       if (result) {
-        session.alwaysOnSkillInstructions =
-          alwaysOnSkillInstructions ?? previous?.alwaysOnSkillInstructions;
+        session.alwaysOnSkillInstructions = resolvedAlwaysOnSkillInstructions;
         const liveConfigOptions = result.configOptions as
           | SessionConfigOption[]
           | undefined;
@@ -2166,6 +2177,7 @@ export class SessionService {
       // permanent disconnect (archive, delete, fresh session) may drop them.
       this.d.adapterStore.removeAdapter(taskRunId);
       this.d.removePersistedConfigOptions(taskRunId);
+      this.d.removePersistedAlwaysOnSkillInstructions?.(taskRunId);
     }
   }
 
@@ -2404,6 +2416,12 @@ export class SessionService {
     session.executionMode = executionMode;
     session.reasoningLevel = reasoningLevel;
     session.alwaysOnSkillInstructions = alwaysOnSkillInstructions;
+    if (alwaysOnSkillInstructions) {
+      this.d.setPersistedAlwaysOnSkillInstructions?.(
+        taskRun.id,
+        alwaysOnSkillInstructions,
+      );
+    }
 
     // An imported CLI session had its history replayed during agent.start;
     // the replay is already in the local run log, so load it for the UI.

packages/ui/src/features/sessions/sessionConfigStore.test.ts

@@ -0,0 +1,49 @@
+import {
+  flushRendererStateWrites,
+  registerRendererStateStorage,
+} from "@posthog/ui/shell/rendererStorage";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import {
+  getPersistedAlwaysOnSkillInstructions,
+  removePersistedAlwaysOnSkillInstructions,
+  setPersistedAlwaysOnSkillInstructions,
+  useSessionConfigStore,
+} from "./sessionConfigStore";
+
+const getItem = vi.fn();
+const setItem = vi.fn();
+const removeItem = vi.fn();
+
+registerRendererStateStorage({ getItem, setItem, removeItem });
+
+describe("sessionConfigStore always-on skill instructions", () => {
+  beforeEach(async () => {
+    await flushRendererStateWrites();
+    getItem.mockReset();
+    setItem.mockReset();
+    removeItem.mockReset();
+    getItem.mockResolvedValue(null);
+    setItem.mockResolvedValue(undefined);
+    removeItem.mockResolvedValue(undefined);
+    useSessionConfigStore.setState({
+      configsByRunId: {},
+      alwaysOnSkillInstructionsByRunId: {},
+    });
+  });
+
+  it("persists instructions by task run until they are removed", async () => {
+    setPersistedAlwaysOnSkillInstructions("run-1", "Follow this skill");
+
+    expect(getPersistedAlwaysOnSkillInstructions("run-1")).toBe(
+      "Follow this skill",
+    );
+    await flushRendererStateWrites();
+    expect(setItem).toHaveBeenCalledWith(
+      "session-config-storage",
+      expect.stringContaining("Follow this skill"),
+    );
+
+    removePersistedAlwaysOnSkillInstructions("run-1");
+    expect(getPersistedAlwaysOnSkillInstructions("run-1")).toBeUndefined();
+  });
+});

packages/ui/src/features/sessions/sessionConfigStore.ts

@@ -4,8 +4,8 @@ import { create } from "zustand";
 import { persist } from "zustand/middleware";
 
 interface SessionConfigState {
-  /** Map of taskRunId -> persisted config options */
   configsByRunId: Record<string, SessionConfigOption[]>;
+  alwaysOnSkillInstructionsByRunId: Record<string, string>;
 }
 
 interface SessionConfigActions {
@@ -15,6 +15,12 @@ interface SessionConfigActions {
   getConfigOptions: (taskRunId: string) => SessionConfigOption[] | undefined;
   /** Remove config options for a task run */
   removeConfigOptions: (taskRunId: string) => void;
+  setAlwaysOnSkillInstructions: (
+    taskRunId: string,
+    instructions: string,
+  ) => void;
+  getAlwaysOnSkillInstructions: (taskRunId: string) => string | undefined;
+  removeAlwaysOnSkillInstructions: (taskRunId: string) => void;
 }
 
 type SessionConfigStore = SessionConfigState & SessionConfigActions;
@@ -23,6 +29,7 @@ export const useSessionConfigStore = create<SessionConfigStore>()(
   persist(
     (set, get) => ({
       configsByRunId: {},
+      alwaysOnSkillInstructionsByRunId: {},
 
       setConfigOptions: (taskRunId, options) =>
         set((state) => ({
@@ -36,11 +43,30 @@ export const useSessionConfigStore = create<SessionConfigStore>()(
           const { [taskRunId]: _removed, ...rest } = state.configsByRunId;
           return { configsByRunId: rest };
         }),
+      setAlwaysOnSkillInstructions: (taskRunId, instructions) =>
+        set((state) => ({
+          alwaysOnSkillInstructionsByRunId: {
+            ...state.alwaysOnSkillInstructionsByRunId,
+            [taskRunId]: instructions,
+          },
+        })),
+      getAlwaysOnSkillInstructions: (taskRunId) =>
+        get().alwaysOnSkillInstructionsByRunId[taskRunId],
+      removeAlwaysOnSkillInstructions: (taskRunId) =>
+        set((state) => {
+          const { [taskRunId]: _removed, ...rest } =
+            state.alwaysOnSkillInstructionsByRunId;
+          return { alwaysOnSkillInstructionsByRunId: rest };
+        }),
     }),
     {
       name: "session-config-storage",
       storage: electronStorage,
-      partialize: (state) => ({ configsByRunId: state.configsByRunId }),
+      partialize: (state) => ({
+        configsByRunId: state.configsByRunId,
+        alwaysOnSkillInstructionsByRunId:
+          state.alwaysOnSkillInstructionsByRunId,
+      }),
     },
   ),
 );
@@ -64,3 +90,26 @@ export function setPersistedConfigOptions(
 export function removePersistedConfigOptions(taskRunId: string): void {
   useSessionConfigStore.getState().removeConfigOptions(taskRunId);
 }
+
+export function getPersistedAlwaysOnSkillInstructions(
+  taskRunId: string,
+): string | undefined {
+  return useSessionConfigStore
+    .getState()
+    .getAlwaysOnSkillInstructions(taskRunId);
+}
+
+export function setPersistedAlwaysOnSkillInstructions(
+  taskRunId: string,
+  instructions: string,
+): void {
+  useSessionConfigStore
+    .getState()
+    .setAlwaysOnSkillInstructions(taskRunId, instructions);
+}
+
+export function removePersistedAlwaysOnSkillInstructions(
+  taskRunId: string,
+): void {
+  useSessionConfigStore.getState().removeAlwaysOnSkillInstructions(taskRunId);
+}

packages/ui/src/features/sessions/sessionServiceHost.ts

@@ -31,8 +31,11 @@ import { NotificationBus } from "@posthog/ui/features/notifications/notification
 import { SpeechNotifier } from "@posthog/ui/features/notifications/speechNotifier";
 import { useSessionAdapterStore } from "@posthog/ui/features/sessions/sessionAdapterStore";
 import {
+  getPersistedAlwaysOnSkillInstructions,
   getPersistedConfigOptions,
+  removePersistedAlwaysOnSkillInstructions,
   removePersistedConfigOptions,
+  setPersistedAlwaysOnSkillInstructions,
   setPersistedConfigOptions,
 } from "@posthog/ui/features/sessions/sessionConfigStore";
 import { sessionStoreSetters } from "@posthog/ui/features/sessions/sessionStore";
@@ -117,6 +120,9 @@ function buildSessionServiceDeps(): SessionServiceDeps {
       getPersistedConfigOptions(taskRunId) ?? undefined,
     setPersistedConfigOptions,
     removePersistedConfigOptions,
+    getPersistedAlwaysOnSkillInstructions,
+    setPersistedAlwaysOnSkillInstructions,
+    removePersistedAlwaysOnSkillInstructions,
     adapterStore: {
       getAdapter: (taskRunId) =>
         useSessionAdapterStore.getState().getAdapter(taskRunId),

packages/ui/src/features/task-detail/hooks/useTaskCreation.ts

@@ -322,30 +322,36 @@ export function useTaskCreation({
       let alwaysOnSkills: AlwaysOnSkillRef[] = settings.alwaysOnSkills;
       let alwaysOnSkillInstructions: string | undefined;
       while (alwaysOnSkills.length > 0) {
-        try {
-          alwaysOnSkillInstructions =
-            await hostClient.skills.renderAlwaysOn.query(alwaysOnSkills);
+        const rendered =
+          await hostClient.skills.renderAlwaysOn.query(alwaysOnSkills);
+        alwaysOnSkillInstructions = rendered.instructions;
+        if (rendered.failures.length === 0) {
           break;
-        } catch (error) {
-          const action = await useAlwaysOnSkillsFailureStore
-            .getState()
-            .confirm(
-              error instanceof Error ? error.message : String(error),
-              alwaysOnSkills,
-            );
-          if (action === "retry") continue;
-          if (action === "cancel") {
-            setIsCreatingTask(false);
-            return false;
-          }
-          if (action === "disable") {
-            for (const skill of alwaysOnSkills) {
-              useSettingsStore.getState().setSkillAlwaysOn(skill, false);
-            }
+        }
+        const failedSkills = rendered.failures.map(({ skill }) => skill);
+        const action = await useAlwaysOnSkillsFailureStore
+          .getState()
+          .confirm(
+            rendered.failures.map(({ error }) => error).join("\n"),
+            failedSkills,
+          );
+        if (action === "retry") continue;
+        if (action === "cancel") {
+          setIsCreatingTask(false);
+          return false;
+        }
+        if (action === "disable") {
+          for (const skill of failedSkills) {
+            useSettingsStore.getState().setSkillAlwaysOn(skill, false);
           }
-          alwaysOnSkills = [];
-          alwaysOnSkillInstructions = undefined;
         }
+        const failedKeys = new Set(
+          failedSkills.map((skill) => `${skill.source}:${skill.path}`),
+        );
+        alwaysOnSkills = alwaysOnSkills.filter(
+          (skill) => !failedKeys.has(`${skill.source}:${skill.path}`),
+        );
+        break;
       }
 
       const shouldShowPendingView = !onTaskCreated && !!plainPromptText;

packages/workspace-server/src/services/skills/schemas.ts

@@ -139,7 +139,15 @@ export const resolveSkillDependenciesInput = z.array(bundleLocalSkillInput);
 export const resolveSkillDependenciesOutput = z.array(bundleLocalSkillInput);
 
 export const renderAlwaysOnSkillsInput = z.array(bundleLocalSkillInput);
-export const renderAlwaysOnSkillsOutput = z.string();
+export const renderAlwaysOnSkillsOutput = z.object({
+  instructions: z.string().optional(),
+  failures: z.array(
+    z.object({
+      skill: bundleLocalSkillInput,
+      error: z.string(),
+    }),
+  ),
+});
 
 export type BundleLocalSkillInput = z.infer<typeof bundleLocalSkillInput>;
 export type BundleLocalSkillOutput = z.infer<typeof bundleLocalSkillOutput>;

packages/workspace-server/src/services/skills/skills.test.ts

@@ -759,10 +759,28 @@ describe("renderAlwaysOnSkillInstructions", () => {
       { name: "first", source: "repo", path: first },
     ]);
 
-    expect(rendered.indexOf("## first")).toBeLessThan(
-      rendered.indexOf("## second"),
+    expect(rendered.instructions?.indexOf("## first")).toBeLessThan(
+      rendered.instructions?.indexOf("## second") ?? -1,
     );
-    expect(rendered).toContain(`Installed at: ${first}`);
-    expect(rendered).not.toContain("description: about first");
+    expect(rendered.instructions).toContain(`Installed at: ${first}`);
+    expect(rendered.instructions).not.toContain("description: about first");
+    expect(rendered.failures).toEqual([]);
+  });
+
+  it("renders readable skills and reports unreadable skills separately", async () => {
+    const readable = await createSkill(repoSkillsDir, "readable");
+
+    const rendered = await makeService().renderAlwaysOnSkillInstructions([
+      {
+        name: "missing",
+        source: "repo",
+        path: path.join(repoSkillsDir, "missing"),
+      },
+      { name: "readable", source: "repo", path: readable },
+    ]);
+
+    expect(rendered.instructions).toContain("## readable");
+    expect(rendered.failures).toHaveLength(1);
+    expect(rendered.failures[0]?.skill.name).toBe("missing");
   });
 });