PaperClipAI/server/src/__tests__/plugin-host-services-span.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";
import { createHostClientHandlers } from "../../../packages/plugins/sdk/src/host-client-factory.js";
import type { WorkerHostCallContext } from "../../../packages/plugins/sdk/src/protocol.js";
import { SANDBOX_STARTUP_SPAN_ATTRS as A } from "@paperclipai/adapter-utils/acpx-engine/startup-timing";
import {
  buildHostServices,
  clampProviderSpanAttributes,
  parseTraceparent,
  recordWorkerProviderSpan,
} from "../services/plugin-host-services.js";

// Capture every span the host trust boundary hands to the real tracer.
const mockRecordSpan = vi.hoisted(() => vi.fn());

vi.mock("../instrumentation.js", () => ({
  recordProviderPluginSpan: mockRecordSpan,
  traceparentFromContextToken: () => undefined,
}));

function createEventBusStub() {
  return {
    forPlugin() {
      return { emit: vi.fn(), subscribe: vi.fn() };
    },
  } as never;
}

// A well-formed W3C traceparent (the host mints it; the handler validates it).
const VALID_TRACEPARENT = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01";

function servicesFor() {
  return buildHostServices({} as never, "plugin-record-id", "daytona", createEventBusStub());
}

function handlersFor(capabilities: readonly string[]) {
  return createHostClientHandlers({
    pluginId: "daytona",
    capabilities: capabilities as never,
    services: servicesFor(),
  });
}

describe("plugin provider span host handler", () => {
  beforeEach(() => {
    mockRecordSpan.mockReset();
  });

  it("records a span with the clamped name and the allowlisted attributes", async () => {
    const services = servicesFor();
    await services.tracer.record(
      {
        name: "pack",
        attributes: {
          [A.provider]: "daytona",
          [A.packWallMs]: 12,
        },
      },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );

    expect(mockRecordSpan).toHaveBeenCalledTimes(1);
    const call = mockRecordSpan.mock.calls[0]![0] as {
      name: string;
      parent: { traceId: string; spanId: string; traceFlags: number };
      attributes: Record<string, unknown>;
    };
    expect(call.name).toBe("sandbox.daytona.pack");
    expect(call.parent).toEqual({
      traceId: "0af7651916cd43dd8448eb211c80319c",
      spanId: "b7ad6b7169203331",
      traceFlags: 1,
    });
    expect(call.attributes[A.provider]).toBe("daytona");
    expect(call.attributes[A.packWallMs]).toBe(12);
  });

  it("drops every forbidden attribute before the span reaches the tracer", async () => {
    const services = servicesFor();
    await services.tracer.record(
      {
        name: "transfer",
        attributes: {
          [A.transferGuardCount]: 2,
          // Forbidden fields that must never ride a span.
          [A.execCommand]: "bash",
          command: "rm -rf /",
          args: "--force",
          stdout: "secret output",
          stderr: "secret error",
          path: "/etc/passwd",
          extra: "leak",
        },
      },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );

    expect(mockRecordSpan).toHaveBeenCalledTimes(1);
    const attributes = (mockRecordSpan.mock.calls[0]![0] as { attributes: Record<string, unknown> })
      .attributes;
    // Only the allowlisted key survives; every forbidden key is dropped.
    expect(attributes).toEqual({ [A.transferGuardCount]: 2 });
    for (const forbidden of [A.execCommand, "command", "args", "stdout", "stderr", "path", "extra"]) {
      expect(attributes).not.toHaveProperty(forbidden);
    }
  });

  it("drops a status message (it could carry standard-stream text)", async () => {
    const services = servicesFor();
    await services.tracer.record(
      { name: "transfer", status: { code: 2, message: "secret error text" } },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    const call = mockRecordSpan.mock.calls[0]![0] as { status?: { code: number; message?: string } };
    expect(call.status).toEqual({ code: 2 });
    expect(call.status).not.toHaveProperty("message");
  });

  it("clamps an unknown span name to sandbox.daytona.other", async () => {
    const services = servicesFor();
    await services.tracer.record(
      { name: "rm -rf / --no-preserve-root" },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    expect((mockRecordSpan.mock.calls[0]![0] as { name: string }).name).toBe(
      "sandbox.daytona.other",
    );
  });

  it("admits each per-round-trip span name to sandbox.daytona.<name>", async () => {
    const services = servicesFor();
    for (const name of [
      "ensureDirectory",
      "checkSymlinkEscape",
      "promote",
      "extractTarball",
      "postUploadCommand",
    ]) {
      await services.tracer.record(
        { name },
        { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
      );
    }
    const recorded = mockRecordSpan.mock.calls.map((c) => (c[0] as { name: string }).name);
    expect(recorded).toEqual([
      "sandbox.daytona.ensureDirectory",
      "sandbox.daytona.checkSymlinkEscape",
      "sandbox.daytona.promote",
      "sandbox.daytona.extractTarball",
      "sandbox.daytona.postUploadCommand",
    ]);
  });

  it("admits the session open and close span names", async () => {
    const services = servicesFor();
    for (const name of ["session.open", "session.close"]) {
      await services.tracer.record(
        { name },
        { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
      );
    }
    const recorded = mockRecordSpan.mock.calls.map((c) => (c[0] as { name: string }).name);
    expect(recorded).toEqual([
      "sandbox.daytona.session.open",
      "sandbox.daytona.session.close",
    ]);
  });

  it("forwards a valid start-time and end-time pair to the recorder", async () => {
    const services = servicesFor();
    const startTimeMs = Date.now() - 4500;
    const endTimeMs = startTimeMs + 4500;
    await services.tracer.record(
      { name: "ensureDirectory", startTimeMs, endTimeMs },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    const call = mockRecordSpan.mock.calls[0]![0] as {
      startTimeMs?: number;
      endTimeMs?: number;
    };
    expect(call.startTimeMs).toBe(startTimeMs);
    expect(call.endTimeMs).toBe(endTimeMs);
  });

  it("accepts a pair whose end is a small skew ahead of the host clock", async () => {
    const services = servicesFor();
    // The worker clock leads the host clock by a few seconds. This small skew
    // is within the allowed bound, so the host keeps the native width.
    const startTimeMs = Date.now() + 5000;
    const endTimeMs = startTimeMs + 1000;
    await services.tracer.record(
      { name: "ensureDirectory", startTimeMs, endTimeMs },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    const call = mockRecordSpan.mock.calls[0]![0] as {
      startTimeMs?: number;
      endTimeMs?: number;
    };
    expect(call.startTimeMs).toBe(startTimeMs);
    expect(call.endTimeMs).toBe(endTimeMs);
  });

  it("drops an invalid timestamp pair so the synchronous path runs", async () => {
    const services = servicesFor();
    const now = Date.now();
    const invalidPairs: Array<{ startTimeMs?: unknown; endTimeMs?: unknown; why: string }> = [
      { startTimeMs: now, endTimeMs: now - 1000, why: "reversed order" },
      { startTimeMs: Number.NaN, endTimeMs: now, why: "non-finite start" },
      { startTimeMs: now, endTimeMs: Number.POSITIVE_INFINITY, why: "non-finite end" },
      { startTimeMs: now, endTimeMs: now + 11 * 60 * 1000, why: "over-ceiling duration" },
      { startTimeMs: now - 2 * 60 * 60 * 1000, endTimeMs: now - 2 * 60 * 60 * 1000 + 10, why: "over-age start" },
      { startTimeMs: now, endTimeMs: now + 2 * 60 * 1000, why: "end far in the future" },
      { startTimeMs: now + 5 * 60 * 1000, endTimeMs: now + 5 * 60 * 1000 + 10, why: "start and end in the future" },
    ];
    for (const pair of invalidPairs) {
      mockRecordSpan.mockReset();
      await services.tracer.record(
        { name: "ensureDirectory", startTimeMs: pair.startTimeMs, endTimeMs: pair.endTimeMs } as never,
        { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
      );
      // The span still records (the synchronous path), but without a timestamp.
      const call = mockRecordSpan.mock.calls[0]![0] as {
        startTimeMs?: number;
        endTimeMs?: number;
      };
      expect(call.startTimeMs, pair.why).toBeUndefined();
      expect(call.endTimeMs, pair.why).toBeUndefined();
    }
  });

  it("records the synchronous path when the timestamp pair is absent", async () => {
    const services = servicesFor();
    await services.tracer.record(
      { name: "pack" },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    const call = mockRecordSpan.mock.calls[0]![0] as {
      startTimeMs?: number;
      endTimeMs?: number;
    };
    expect(call.startTimeMs).toBeUndefined();
    expect(call.endTimeMs).toBeUndefined();
  });

  it("rejects a malformed traceparent — no span is recorded", async () => {
    const services = servicesFor();
    for (const bad of [
      undefined,
      "not-a-traceparent",
      "00-xyz-b7ad6b7169203331-01",
      "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331", // missing flags
      "00-00000000000000000000000000000000-b7ad6b7169203331-01", // all-zero trace id
    ]) {
      await services.tracer.record(
        { name: "pack" },
        { traceparent: bad } as WorkerHostCallContext,
      );
    }
    expect(mockRecordSpan).not.toHaveBeenCalled();
  });

  it("rejects a span from a plugin that lacks the environment-driver capability", async () => {
    const handlers = handlersFor([]);
    await expect(
      handlers["span.record"](
        { name: "pack" },
        { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
      ),
    ).rejects.toThrow(/capabilit/i);
    expect(mockRecordSpan).not.toHaveBeenCalled();
  });

  it("admits a span from a plugin that holds the environment-driver capability", async () => {
    const handlers = handlersFor(["environment.drivers.register"]);
    await handlers["span.record"](
      { name: "pack", attributes: { [A.provider]: "daytona" } },
      { traceparent: VALID_TRACEPARENT } as WorkerHostCallContext,
    );
    expect(mockRecordSpan).toHaveBeenCalledTimes(1);
  });
});

describe("parseTraceparent", () => {
  it("accepts a well-formed traceparent and returns the parts", () => {
    expect(parseTraceparent(VALID_TRACEPARENT)).toEqual({
      traceId: "0af7651916cd43dd8448eb211c80319c",
      spanId: "b7ad6b7169203331",
      traceFlags: 1,
    });
  });

  it("rejects malformed, all-zero, and forbidden-version traceparents", () => {
    expect(parseTraceparent(undefined)).toBeNull();
    expect(parseTraceparent("garbage")).toBeNull();
    expect(parseTraceparent("00-0af7651916cd43dd8448eb211c80319c-0000000000000000-01")).toBeNull();
    expect(parseTraceparent("ff-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01")).toBeNull();
  });
});

describe("clampProviderSpanAttributes", () => {
  it("keeps only allowlisted keys and normalizes the provider family", () => {
    expect(
      clampProviderSpanAttributes({
        [A.provider]: "some-operator-key",
        [A.packWallMs]: 7,
        [A.transferWallMs]: Number.NaN,
        [A.execCommand]: "bash",
      }),
    ).toEqual({
      // An unknown provider key maps to `plugin`, never the raw key.
      [A.provider]: "plugin",
      [A.packWallMs]: 7,
      // A non-finite number yields no attribute; `exec.command` is not allowed.
    });
  });

  it("keeps the transfer direction for each closed-set value", () => {
    expect(clampProviderSpanAttributes({ [A.transferDirection]: "inbound" })).toEqual({
      [A.transferDirection]: "inbound",
    });
    expect(clampProviderSpanAttributes({ [A.transferDirection]: "outbound" })).toEqual({
      [A.transferDirection]: "outbound",
    });
  });

  it("drops a transfer direction outside the closed set or of the wrong type", () => {
    // A free-form string, an empty string, and a non-string all yield no
    // attribute, so the direction stays bounded and low-cardinality.
    expect(clampProviderSpanAttributes({ [A.transferDirection]: "sideways" })).toEqual({});
    expect(clampProviderSpanAttributes({ [A.transferDirection]: "" })).toEqual({});
    expect(clampProviderSpanAttributes({ [A.transferDirection]: 1 })).toEqual({});
  });
});

describe("nested worker timing receipts",()=>{
 it("keeps lease verification beneath its worker and clamps provider state", () => {
  mockRecordSpan.mockReset();
  const rootParent = { traceId: "0af7651916cd43dd8448eb211c80319c", spanId: "1111111111111111", traceFlags: 1 };
  const verifyParent = { ...rootParent, spanId: "2222222222222222" };
  mockRecordSpan.mockReturnValueOnce(rootParent).mockReturnValueOnce(verifyParent);
  const now = Date.now();
  recordWorkerProviderSpan({ name: "plugin.worker", startTimeMs: now - 100, endTimeMs: now, children: [
    { name: "lease.verify", startTimeMs: now - 90, endTimeMs: now - 10, children: [
      { name: "lease.sentinel", startTimeMs: now - 80, endTimeMs: now - 20,
        attributes: { "paperclip.sandbox.startup.state": "archived", token: "private-token" } },
    ] },
  ] }, { traceparent: VALID_TRACEPARENT });
  expect(mockRecordSpan.mock.calls[1][0]).toMatchObject({ name: "sandbox.daytona.lease.verify", parent: rootParent });
  expect(mockRecordSpan.mock.calls[2][0]).toMatchObject({ name: "sandbox.daytona.lease.sentinel", parent: verifyParent,
    startTimeMs: now - 80, endTimeMs: now - 20, attributes: { "paperclip.sandbox.startup.state": "archived" } });
  expect(JSON.stringify(mockRecordSpan.mock.calls)).not.toContain("private-token");
  expect(clampProviderSpanAttributes({ "paperclip.sandbox.startup.state": "private-state" })).toEqual({});
 });

 it("admits a workspace failure span with only a closed failure reason", () => {
  mockRecordSpan.mockReset();
  const now = Date.now();
  recordWorkerProviderSpan({ name: "plugin.workspace.failure", startTimeMs: now - 1, endTimeMs: now, status: { code: 2 },
    attributes: { "plugin.failure_reason": "storage_grant_unavailable", message: "Project mount grants unavailable" } }, { traceparent: VALID_TRACEPARENT });
  expect(mockRecordSpan.mock.calls[0][0]).toMatchObject({ name: "plugin.workspace.failure", attributes: { "plugin.failure_reason": "storage_grant_unavailable" } });
  expect(JSON.stringify(mockRecordSpan.mock.calls)).not.toContain("mount grants");
  expect(clampProviderSpanAttributes({ "plugin.failure_reason": "unknown" })).toEqual({ "plugin.failure_reason": "unknown" });
  expect(clampProviderSpanAttributes({ "plugin.failure_reason": "token=private" })).toEqual({});
  expect(clampProviderSpanAttributes({ "plugin.failure_reason": 1 })).toEqual({});
 });

 it("admits every closed lease phase and rejects dynamic phase names", () => {
  mockRecordSpan.mockReset();
  const names = ["lease.wait_for_teardown", "lease.lookup", "lease.refresh", "lease.channel_cleanup",
    "lease.start", "lease.recover", "lease.verify", "lease.workspace", "lease.sentinel", "lease.shell",
    "lease.create", "lease.expiry", "lease.cleanup", "lease.private-resource-id"];
  for (const name of names) recordWorkerProviderSpan({ name }, { traceparent: VALID_TRACEPARENT });
  expect(mockRecordSpan.mock.calls.map(call => call[0].name)).toEqual([
    ...names.slice(0, -1).map(name => `sandbox.daytona.${name}`), "sandbox.daytona.other",
  ]);
 });

 it("uses host-generated parents, clamps content and bounds child clocks",()=>{
  mockRecordSpan.mockReset();
  const childParent={traceId:"0af7651916cd43dd8448eb211c80319c",spanId:"1111111111111111",traceFlags:1};
  mockRecordSpan.mockReturnValue(childParent);
  const now=Date.now();
  recordWorkerProviderSpan({name:"plugin.worker",startTimeMs:now-100,endTimeMs:now,children:[
   {name:"plugin.remote.operation",startTimeMs:now-90,endTimeMs:now-10,attributes:{"plugin.operation":"script.settings",secret:"private"}},
   {name:"plugin.runtime.verify",startTimeMs:now-200,endTimeMs:now},
  ]},{traceparent:VALID_TRACEPARENT} as WorkerHostCallContext);
  expect(mockRecordSpan).toHaveBeenCalledTimes(2);
  expect(mockRecordSpan.mock.calls[1][0]).toMatchObject({parent:childParent,attributes:{"plugin.operation":"script.settings"}});
  expect(JSON.stringify(mockRecordSpan.mock.calls)).not.toContain("private");
 });
});

it("allows a workspace provider timing receipt without environment-driver authority",async()=>{
 mockRecordSpan.mockReset();
 await handlersFor(["workspaces.manage"])["span.record"]({name:"plugin.worker"},{traceparent:VALID_TRACEPARENT} as WorkerHostCallContext);
 expect(mockRecordSpan).toHaveBeenCalledTimes(1);
});

it("retains short children of a multi-hour worker invocation and shares its aggregate budget",()=>{
 mockRecordSpan.mockReset();mockRecordSpan.mockReturnValue({traceId:"0af7651916cd43dd8448eb211c80319c",spanId:"1111111111111111",traceFlags:1});
 const now=Date.now(),context={traceparent:VALID_TRACEPARENT,providerSpanBudget:{remaining:3}};
 recordWorkerProviderSpan({name:"plugin.worker",startTimeMs:now-7200000,endTimeMs:now,children:[
  {name:"pack",startTimeMs:now-7199990,endTimeMs:now-7199980}
 ]},context);
 recordWorkerProviderSpan({name:"pack"},context);recordWorkerProviderSpan({name:"pack"},context);
 expect(mockRecordSpan).toHaveBeenCalledTimes(3);
 expect(mockRecordSpan.mock.calls[1][0]).toMatchObject({name:"sandbox.daytona.pack",startTimeMs:now-7199990});
});


it("links bounded earlier remote lifecycles to their host-owned observer",()=>{
 mockRecordSpan.mockReset();mockRecordSpan.mockReturnValue({traceId:"0af7651916cd43dd8448eb211c80319c",spanId:"1111111111111111",traceFlags:1});
 const now=Date.now();
 recordWorkerProviderSpan({name:"plugin.worker",startTimeMs:now-10,endTimeMs:now,children:[
  {name:"plugin.remote.lifecycle",relationship:"linked",startTimeMs:now-1000,endTimeMs:now-100,children:[
   {name:"plugin.remote.operation",startTimeMs:now-900,endTimeMs:now-200}
  ]},
  {name:"plugin.remote.operation",relationship:"linked",startTimeMs:now-1000,endTimeMs:now-100},
  {name:"plugin.remote.lifecycle",relationship:"linked",startTimeMs:now-700000,endTimeMs:now-100},
 ]},{traceparent:VALID_TRACEPARENT});
 expect(mockRecordSpan).toHaveBeenCalledTimes(3);
 expect(mockRecordSpan.mock.calls[1][0]).toMatchObject({name:"plugin.remote.lifecycle",linked:true,startTimeMs:now-1000});
 expect(mockRecordSpan.mock.calls[2][0].linked).toBeUndefined();
});