Fix plugin company-scoped config resolution

server/src/__tests__/plugin-host-services-session-events.test.ts

@@ -0,0 +1,129 @@
+import { randomUUID } from "node:crypto";
+import { afterAll, afterEach, beforeAll, describe, expect, it, vi } from "vitest";
+import { agentTaskSessions, agents, companies, createDb } from "@paperclipai/db";
+import {
+  getEmbeddedPostgresTestSupport,
+  startEmbeddedPostgresTestDatabase,
+} from "./helpers/embedded-postgres.js";
+
+let liveEventListener: ((event: {
+  type: string;
+  payload?: Record<string, unknown>;
+}) => void) | null = null;
+
+vi.mock("../services/heartbeat.js", () => ({
+  heartbeatService: () => ({
+    wakeup: vi.fn(async () => ({ id: "run-1" })),
+  }),
+}));
+
+vi.mock("../services/live-events.js", () => ({
+  subscribeCompanyLiveEvents: vi.fn((_companyId: string, listener: typeof liveEventListener) => {
+    liveEventListener = listener;
+    return () => {
+      liveEventListener = null;
+    };
+  }),
+}));
+
+const { buildHostServices } = await import("../services/plugin-host-services.js");
+
+const embeddedPostgresSupport = await getEmbeddedPostgresTestSupport();
+const describeEmbeddedPostgres = embeddedPostgresSupport.supported ? describe : describe.skip;
+
+describeEmbeddedPostgres("plugin host services session events", () => {
+  let db!: ReturnType<typeof createDb>;
+  let tempDb: Awaited<ReturnType<typeof startEmbeddedPostgresTestDatabase>> | null = null;
+
+  beforeAll(async () => {
+    tempDb = await startEmbeddedPostgresTestDatabase("paperclip-plugin-host-services-");
+    db = createDb(tempDb.connectionString);
+  }, 20_000);
+
+  afterEach(async () => {
+    liveEventListener = null;
+    await db.delete(agentTaskSessions);
+    await db.delete(agents);
+    await db.delete(companies);
+  });
+
+  afterAll(async () => {
+    await tempDb?.cleanup();
+  });
+
+  function createEventBusStub() {
+    return {
+      forPlugin() {
+        return {
+          emit: async () => {},
+          subscribe: () => {},
+          clear: () => {},
+        };
+      },
+    } as any;
+  }
+
+  async function seedCompanyAndAgent() {
+    const companyId = randomUUID();
+    const agentId = randomUUID();
+    await db.insert(companies).values({
+      id: companyId,
+      name: "Paperclip",
+      issuePrefix: `T${companyId.replace(/-/g, "").slice(0, 6).toUpperCase()}`,
+      requireBoardApprovalForNewAgents: false,
+    });
+    await db.insert(agents).values({
+      id: agentId,
+      companyId,
+      name: "Engineer",
+      role: "engineer",
+      status: "idle",
+      adapterType: "process",
+      adapterConfig: { command: "true" },
+      runtimeConfig: {},
+      permissions: {},
+    });
+    return { companyId, agentId };
+  }
+
+  it("forwards companyId on live agent session notifications", async () => {
+    const { companyId, agentId } = await seedCompanyAndAgent();
+    const notifyWorker = vi.fn();
+    const services = buildHostServices(
+      db,
+      "plugin-record-id",
+      "paperclip.missions",
+      createEventBusStub(),
+      notifyWorker,
+    );
+
+    const session = await services.agentSessions.create({ agentId, companyId });
+    const sendResult = await services.agentSessions.sendMessage({
+      sessionId: session.sessionId,
+      companyId,
+      prompt: "hello",
+    });
+
+    liveEventListener?.({
+      type: "heartbeat.run.log",
+      payload: {
+        runId: sendResult.runId,
+        seq: 1,
+        stream: "stdout",
+        chunk: "hello from run",
+      },
+    });
+
+    expect(notifyWorker).toHaveBeenCalledWith(
+      "agents.sessions.event",
+      expect.objectContaining({
+        sessionId: session.sessionId,
+        runId: sendResult.runId,
+        companyId,
+        eventType: "chunk",
+      }),
+    );
+
+    services.dispose();
+  });
+});

server/src/__tests__/plugin-routes-authz.test.ts

@@ -86,7 +86,9 @@ async function createApp(
     next();
   });
   app.use("/api", pluginRoutes(
-    (routeOverrides.db ?? {}) as never,
+    (routeOverrides.db ?? {
+      transaction: async <T>(callback: (tx: unknown) => Promise<T>) => callback({}),
+    }) as never,
     loader as never,
     routeOverrides.jobDeps as never,
     undefined,
@@ -363,7 +365,9 @@ describe.sequential("plugin install and upgrade authz", () => {
     const { app } = await createApp(boardActor({
       isInstanceAdmin: true,
       companyIds: [companyA],
-    }));
+    }), {}, {
+      bridgeDeps: { workerManager: mockWorkerManager },
+    });
 
     const res = await request(app)
       .post(`/api/plugins/${pluginId}/config`)
@@ -382,6 +386,7 @@ describe.sequential("plugin install and upgrade authz", () => {
         secretId: "77777777-7777-4777-8777-777777777777",
         configPath: "$",
       }],
+      { db: {} },
     );
     expect(mockRegistry.upsertConfig).toHaveBeenCalledWith(
       pluginId,
@@ -392,6 +397,7 @@ describe.sequential("plugin install and upgrade authz", () => {
       },
       companyA,
     );
+    expect(mockWorkerManager.call).not.toHaveBeenCalled();
   }, 20_000);
 
   it("rejects company-scoped plugin config secret refs that do not belong to the selected company", async () => {

server/src/routes/plugins.ts

@@ -79,6 +79,8 @@ import {
 import { secretService } from "../services/secrets.js";
 import { badRequest, forbidden, notFound, unauthorized, unprocessable } from "../errors.js";
 
+type DbTransaction = Parameters<Parameters<Db["transaction"]>[0]>[0];
+
 /** UI slot declaration extracted from plugin manifest */
 type PluginUiSlotDeclaration = NonNullable<NonNullable<PaperclipPluginManifestV1["ui"]>["slots"]>[number];
 /** Launcher declaration extracted from plugin manifest */
@@ -691,6 +693,30 @@ export function pluginRoutes(
     return null;
   }
 
+  async function upsertPluginConfigWithSecretRefs(
+    pluginId: string,
+    configJson: Record<string, unknown>,
+    companyId: string | null,
+    secretRefsByPath: Map<string, Set<string>>,
+  ) {
+    return db.transaction(async (tx) => {
+      const txRegistry = pluginRegistryService(tx as unknown as Db);
+      if (companyId) {
+        const refs = [...secretRefsByPath.entries()].flatMap(([secretId, paths]) =>
+          [...paths].map((configPath) => ({ secretId, configPath })),
+        );
+        await secrets.syncSecretRefsForTarget(
+          companyId,
+          { targetType: "plugin", targetId: pluginId },
+          refs,
+          { db: tx as DbTransaction },
+        );
+      }
+
+      return txRegistry.upsertConfig(pluginId, { configJson }, companyId);
+    });
+  }
+
   function assertPluginBridgeScope(req: Request, companyId: unknown): string | undefined {
     if (companyId === undefined || companyId === null) {
       assertInstanceAdmin(req);
@@ -2201,20 +2227,12 @@ export function pluginRoutes(
         res.status(422).json({ error: "Plugin secret references require companyId" });
         return;
       }
-      if (companyId) {
-        const refs = [...secretRefsByPath.entries()].flatMap(([secretId, paths]) =>
-          [...paths].map((configPath) => ({ secretId, configPath })),
-        );
-        await secrets.syncSecretRefsForTarget(
-          companyId,
-          { targetType: "plugin", targetId: plugin.id },
-          refs,
-        );
-      }
-
-      const result = await registry.upsertConfig(plugin.id, {
-        configJson: body.configJson,
-      }, companyId);
+      const result = await upsertPluginConfigWithSecretRefs(
+        plugin.id,
+        body.configJson,
+        companyId,
+        secretRefsByPath,
+      );
       await logPluginMutationActivity(req, "plugin.config.updated", plugin.id, {
         pluginId: plugin.id,
         pluginKey: plugin.pluginKey,

server/src/services/plugin-host-services.ts

@@ -2647,6 +2647,7 @@ export function buildHostServices(
               notifyWorker("agents.sessions.event", {
                 sessionId: params.sessionId,
                 runId: run.id,
+                companyId,
                 seq: (payload.seq as number) ?? 0,
                 eventType: "chunk",
                 stream: (payload.stream as string) ?? null,
@@ -2659,6 +2660,7 @@ export function buildHostServices(
                 notifyWorker("agents.sessions.event", {
                   sessionId: params.sessionId,
                   runId: run.id,
+                  companyId,
                   seq: 0,
                   eventType: status === "succeeded" ? "done" : "error",
                   stream: "system",
@@ -2670,6 +2672,7 @@ export function buildHostServices(
                 notifyWorker("agents.sessions.event", {
                   sessionId: params.sessionId,
                   runId: run.id,
+                  companyId,
                   seq: 0,
                   eventType: "status",
                   stream: "system",

server/src/services/plugin-registry.ts

@@ -1,4 +1,4 @@
-import { asc, eq, ne, sql, and, isNull } from "drizzle-orm";
+import { asc, eq, ne, sql, and, isNull, or } from "drizzle-orm";
 import type { Db } from "@paperclipai/db";
 import {
   plugins,
@@ -44,13 +44,27 @@ function isPluginKeyConflict(error: unknown): boolean {
   return err.code === "23505" && constraint === "plugins_plugin_key_idx";
 }
 
-function pluginConfigScopeCondition(pluginId: string, companyId?: string | null) {
+function pluginConfigExactScopeCondition(pluginId: string, companyId?: string | null) {
   return and(
     eq(pluginConfig.pluginId, pluginId),
     companyId ? eq(pluginConfig.companyId, companyId) : isNull(pluginConfig.companyId),
   );
 }
 
+function pluginConfigReadScopeCondition(pluginId: string, companyId?: string | null) {
+  if (!companyId) {
+    return pluginConfigExactScopeCondition(pluginId, null);
+  }
+
+  return and(
+    eq(pluginConfig.pluginId, pluginId),
+    or(
+      eq(pluginConfig.companyId, companyId),
+      isNull(pluginConfig.companyId),
+    ),
+  );
+}
+
 // ---------------------------------------------------------------------------
 // Service
 // ---------------------------------------------------------------------------
@@ -292,7 +306,12 @@ export function pluginRegistryService(db: Db) {
       db
         .select()
         .from(pluginConfig)
-        .where(pluginConfigScopeCondition(pluginId, companyId))
+        .where(pluginConfigReadScopeCondition(pluginId, companyId))
+        .orderBy(
+          companyId
+            ? sql`case when ${pluginConfig.companyId} = ${companyId} then 0 else 1 end`
+            : sql`0`,
+        )
         .then((rows) => rows[0] ?? null),
 
     /**
@@ -307,7 +326,7 @@ export function pluginRegistryService(db: Db) {
       const existing = await db
         .select()
         .from(pluginConfig)
-        .where(pluginConfigScopeCondition(pluginId, companyId))
+        .where(pluginConfigExactScopeCondition(pluginId, companyId))
         .then((rows) => rows[0] ?? null);
 
       if (existing) {
@@ -318,7 +337,7 @@ export function pluginRegistryService(db: Db) {
             lastError: null,
             updatedAt: new Date(),
           })
-          .where(pluginConfigScopeCondition(pluginId, companyId))
+          .where(pluginConfigExactScopeCondition(pluginId, companyId))
           .returning()
           .then((rows) => rows[0]);
       }
@@ -345,7 +364,7 @@ export function pluginRegistryService(db: Db) {
       const existing = await db
         .select()
         .from(pluginConfig)
-        .where(pluginConfigScopeCondition(pluginId, companyId))
+        .where(pluginConfigExactScopeCondition(pluginId, companyId))
         .then((rows) => rows[0] ?? null);
 
       if (existing) {
@@ -357,7 +376,7 @@ export function pluginRegistryService(db: Db) {
             lastError: null,
             updatedAt: new Date(),
           })
-          .where(pluginConfigScopeCondition(pluginId, companyId))
+          .where(pluginConfigExactScopeCondition(pluginId, companyId))
           .returning()
           .then((rows) => rows[0]);
       }
@@ -381,7 +400,7 @@ export function pluginRegistryService(db: Db) {
       const rows = await db
         .update(pluginConfig)
         .set({ lastError, updatedAt: new Date() })
-        .where(pluginConfigScopeCondition(pluginId, companyId))
+        .where(pluginConfigExactScopeCondition(pluginId, companyId))
         .returning();
 
       if (rows.length === 0) throw notFound("Plugin config not found");
@@ -392,7 +411,7 @@ export function pluginRegistryService(db: Db) {
     deleteConfig: async (pluginId: string, companyId?: string | null) => {
       const rows = await db
         .delete(pluginConfig)
-        .where(pluginConfigScopeCondition(pluginId, companyId))
+        .where(pluginConfigExactScopeCondition(pluginId, companyId))
         .returning();
 
       return rows[0] ?? null;