dsh-context/tests/host/attribution.spec.ts

// The runtime register() attribution hook (src/host/attribution.ts) against a REAL cordis Context:
// plugin fibers read `ctx.tools` (firing the internal/get waterfall) and call `register()` on a fake tools service.

import assert from 'node:assert/strict'
import path from 'node:path'
import { fileURLToPath } from 'node:url'
import { describe, test } from 'vitest'
import { Context, Service } from '@deepseek-ai/cordis'
import { UNKNOWN_TOOL_SOURCE } from '../../src/shared/types'
import { callerPackageFrom, createToolAttribution, packageNameFrom, type ToolAttribution } from '../../src/host/attribution'

const fileUrl = (file: string) => 'file:///' + file.replace(/\\/g, '/')

/** Absolute path of this spec file (a file inside the dsh-context package). */
const specFile = path.normalize(fileURLToPath(import.meta.url))

/** Symbol cordis registers on its traced proxies (`Symbol.for`, shared
 * registry); reading it yields the wrapped underlying service instance. */
const CORDIS_ORIGINAL = Symbol.for('cordis.original')

/** A tracked tools service shaped like dsh's ToolRuntime: a real Service
 * subclass whose `register` is a PROTOTYPE method — the shape that makes
 * cordis hand out a fresh traced proxy (closing over the reader's context) on every `ctx.tools` property read. */
class TrackedTools extends Service {
  layers = { global: { tools: new Map() } }
  constructor(ctx: Context) { super(ctx, 'tools') }
  register(_definition?: { name?: unknown }) {
    return () => {}
  }
}

describe('createToolAttribution', () => {
  test('falls back to the static chain when nothing was registered at runtime', () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    assert.equal(attribution.ownerOf('read'), '@deepseek-ai/dsh-tool-fs')
    assert.equal(attribution.ownerOf('mcp__github__get_issue'), 'mcp:github')
    assert.equal(attribution.ownerOf('unknown_tool'), undefined)
  })

  test('attributes tools to the plugin that read ctx.tools and registered them', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    let release: () => void = () => {}
    await app.plugin({
      name: 'provider-a',
      apply(sub) {
        const tools = (sub as any).tools
        assert.equal((sub as any).tools, fake, 'a wrapped instance is returned as-is')
        release = tools.register({ name: 'alpha' })
        tools.register({ name: 'beta' })
      },
    })
    assert.equal(attribution.ownerOf('alpha'), 'provider-a')
    assert.equal(attribution.ownerOf('beta'), 'provider-a')
    release()
    assert.equal(attribution.ownerOf('alpha'), undefined, 'the disposer forgets the tool')
    assert.equal(attribution.ownerOf('beta'), 'provider-a')
  })

  test('falls back to the caller stack when no tools read happened', async () => {
    const app = new Context()
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    const attribution = createToolAttribution(app)
    const { registerBlind } = await import('./fixtures/anon-plugin/index.js')
    const tools = app.get('tools', false) as { register(def?: unknown): () => void }
    const release = registerBlind(tools)
    assert.equal(attribution.ownerOf('blind_tool'), 'anon-plugin', 'no reader → the caller stack wins')
    release()
    assert.equal(attribution.ownerOf('blind_tool'), undefined, 'the disposer forgets the tool')
  })

  test('root and self reader slots fall back to the caller stack', async () => {
    const app = new Context()
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    const { registerBlind } = await import('./fixtures/anon-plugin/index.js')
    let attribution!: ToolAttribution
    await app.plugin({
      name: 'dsh-context',
      inject: ['tools'],
      apply(sub) {
        attribution = createToolAttribution(sub)
        registerBlind((sub as any).tools)
      },
    })
    await app.plugin({
      apply(sub) {
        registerBlind((sub as any).tools)
      },
    })
    assert.equal(attribution.ownerOf('blind_tool'), 'anon-plugin', 'self/root slots are never taken as-is')
  })

  test('ignores non-tools reads and keeps the static chain intact', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    app.provide('bench', { value: 1 })
    await app.plugin({
      name: 'probe',
      apply(sub) {
        void (sub as any).bench.value
        void (sub as any).bench.value
      },
    })
    assert.equal(attribution.ownerOf('grep'), '@deepseek-ai/dsh-tool-fs-search')
  })

  test('tolerates falsy and non-object tools values', async () => {
    for (const value of [undefined, 42]) {
      const app = new Context()
      const attribution = createToolAttribution(app)
      app.provide('tools', value)
      await app.plugin({
        name: 'probe',
        apply(sub) {
          void (sub as any).tools
        },
      })
      assert.equal(attribution.ownerOf('x'), undefined)
    }
  })

  test('leaves instances without a callable register untouched', async () => {
    for (const tools of [{ register: 42 }, {}]) {
      const app = new Context()
      const attribution = createToolAttribution(app)
      app.provide('tools', tools)
      await app.plugin({
        name: 'probe',
        apply(sub) {
          void (sub as any).tools
        },
      })
      assert.equal(attribution.ownerOf('whatever'), undefined)
    }
  })

  test('a foreign wrapper carrying a garbage attribution marker is left untouched', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    // A wrapper that copied the peel-marker name with a non-function value:
    // there is no original to recover, so the instance is not re-wrapped.
    const foreign = Object.assign(() => 'not-a-disposer', { attributedOriginal: 42 })
    const fake = { register: foreign }
    app.provide('tools', fake)
    await app.plugin({
      name: 'probe',
      apply(sub) {
        void (sub as any).tools
      },
    })
    assert.equal((app.get('tools', false) as typeof fake).register, foreign, 'the instance was not re-wrapped')
    assert.equal(attribution.ownerOf('whatever'), undefined)
  })

  test('ignores registrations without a string tool name or a disposer', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    app.provide('tools', { register() { return undefined } })
    await app.plugin({
      name: 'weird',
      apply(sub) {
        const tools = (sub as any).tools
        tools.register()
        tools.register({ name: 42 })
        tools.register({ name: 'sticky' })
      },
    })
    assert.equal(attribution.ownerOf('sticky'), 'weird', 'records even without a disposer')
    assert.equal(attribution.ownerOf('42'), undefined, 'non-string tool names are not recorded')
  })

  test('the live record outranks the pinned map; name-derived MCP labels outrank the live record', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    await app.plugin({
      name: 'third-party',
      apply(sub) {
        const tools = (sub as any).tools
        tools.register({ name: 'bash' })
        tools.register({ name: 'mcp__github__x' })
        tools.register({ name: 'brand_new_tool' })
      },
    })
    assert.equal(attribution.ownerOf('bash'), 'third-party', 'a post-boot registration is the truth, even under a first-party name')
    assert.equal(attribution.ownerOf('mcp__github__x'), 'mcp:github', 'the name-derived MCP label names the actual provider')
    assert.equal(attribution.ownerOf('brand_new_tool'), 'third-party', 'live record fills the gap')
  })

  test('a pinned first-party name still resolves when it was never registered at runtime', () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    assert.equal(attribution.ownerOf('bash'), '@deepseek-ai/dsh-tool-bash', 'the pinned map serves boot-time first-party tools')
  })

  test('re-installing the hook re-wraps the instance without stacking wrappers', async () => {
    const app = new Context()
    let calls = 0
    app.provide('tools', { register() { calls++; return () => {} } })
    const first = createToolAttribution(app)
    void first
    // The second hook incarnation (a reload) peels the first wrapper back to
    // the original: a registration records into the NEW hook only, and the original register still runs exactly once.
    const second = createToolAttribution(app)
    await app.plugin({
      name: 'provider-b',
      apply(sub) {
        ;((sub as any).tools as { register(def?: unknown): () => void }).register({ name: 'reload_tool' })
      },
    })
    assert.equal(calls, 1, 'the original register ran once, through one wrapper only')
    assert.equal(second.ownerOf('reload_tool'), 'provider-b', 'the fresh hook records the registration')
  })

  test('unloading the hook plugin restores the original register', async () => {
    const app = new Context()
    const fake = { register(_definition?: unknown) { return () => {} } }
    app.provide('tools', fake)
    const original = fake.register
    let attribution!: ToolAttribution
    const fiber = app.plugin({
      name: 'dsh-context',
      apply(sub) {
        attribution = createToolAttribution(sub)
      },
    })
    await fiber
    assert.notEqual(fake.register, original, 'the hook wrapped the instance at apply')
    await fiber.dispose()
    assert.equal(fake.register, original, 'unload put the original register back on the instance')
    fake.register({ name: 'orphan_tool' })
    assert.equal(attribution.ownerOf('orphan_tool'), undefined, 'no dead-hook recording after unload')
  })

  test('an unloaded older incarnation leaves a newer wrapper installed', async () => {
    const app = new Context()
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    const original = fake.register
    let second!: ToolAttribution
    const older = app.plugin({ name: 'older-hook', apply(sub) { createToolAttribution(sub) } })
    await older
    assert.notEqual(fake.register, original, 'the first hook wrapped the instance')
    const newer = app.plugin({ name: 'newer-hook', apply(sub) { second = createToolAttribution(sub) } })
    await newer
    assert.notEqual(fake.register, original, 'the newer incarnation re-wrapped the instance')
    await older.dispose()
    assert.notEqual(fake.register, original, 'the older disposer did not strip the newer wrapper')
    await app.plugin({
      name: 'provider-c',
      apply(sub) {
        ;((sub as any).tools as { register(def?: unknown): () => void }).register({ name: 'late_tool' })
      },
    })
    assert.equal(second.ownerOf('late_tool'), 'provider-c', 'the surviving hook still records')
    await newer.dispose()
    assert.equal(fake.register, original, 'the newer disposer restores the true original')
  })

  test('unload then re-apply re-wraps the restored original without stacking', async () => {
    const app = new Context()
    let calls = 0
    app.provide('tools', { register() { calls++; return () => {} } })
    let attribution!: ToolAttribution
    const first = app.plugin({ name: 'dsh-context', apply(sub) { createToolAttribution(sub) } })
    await first
    await first.dispose()
    const second = app.plugin({
      name: 'dsh-context',
      apply(sub) {
        attribution = createToolAttribution(sub)
      },
    })
    await second
    await app.plugin({
      name: 'provider-d',
      apply(sub) {
        ;((sub as any).tools as { register(def?: unknown): () => void }).register({ name: 'reloaded_tool' })
      },
    })
    assert.equal(calls, 1, 'the original register ran once, through one wrapper only')
    assert.equal(attribution.ownerOf('reloaded_tool'), 'provider-d', 'the re-applied hook records the registration')
  })

  test('a tracked service is wrapped once on its stable identity and restored on unload', async () => {
    const app = new Context()
    const tools = new TrackedTools(app)
    const original = tools.register
    let attribution!: ToolAttribution
    const fiber = app.plugin({
      name: 'dsh-context',
      apply(sub) {
        attribution = createToolAttribution(sub)
      },
    })
    await fiber
    assert.notEqual(tools.register, original, 'the hook wrapped the tracked instance at apply')
    const wrappedOnce = tools.register
    await app.plugin({
      name: 'agent-scope',
      apply(sub) {
        // Tracked services yield a FRESH traced proxy per property read, both
        // peeling to the same underlying instance (issue #108's premise).
        const first = (sub as unknown as { tools: TrackedTools & Record<symbol, unknown> }).tools
        const second = (sub as unknown as { tools: TrackedTools & Record<symbol, unknown> }).tools
        assert.notEqual(first, second, 'cordis hands out a fresh traced proxy per read')
        assert.equal(first[CORDIS_ORIGINAL], tools, 'the proxy peels to the underlying service')
        assert.equal(second[CORDIS_ORIGINAL], tools)
        assert.equal(tools.register, wrappedOnce, 'repeated reads never re-wrap the stable instance')
        first.register({ name: 'scoped_tool' })
      },
    })
    assert.equal(attribution.ownerOf('scoped_tool'), 'agent-scope', 'the hook still attributes scoped registrations')
    await fiber.dispose()
    assert.equal(tools.register, original, 'unload restored the prototype register of the tracked service')
    const orphan = tools.register({ name: 'orphan_tool' })
    assert.equal(attribution.ownerOf('orphan_tool'), undefined, 'no dead-hook recording after unload')
    orphan()
  })

  test('tags boot-time tools whose providers predate the hook as unknown', () => {
    const app = new Context()
    const fake = {
      register() { return () => {} },
      layers: { global: { tools: { entries: () => [
        ['claim_files', {}],
        ['pending_write', {}],
        ['read', {}],
      ] } } },
    }
    app.provide('tools', fake)
    const attribution = createToolAttribution(app)
    assert.equal(attribution.ownerOf('claim_files'), UNKNOWN_TOOL_SOURCE, 'boot-only tools carry the unknown sentinel')
    assert.equal(attribution.ownerOf('pending_write'), UNKNOWN_TOOL_SOURCE)
    assert.equal(attribution.ownerOf('read'), '@deepseek-ai/dsh-tool-fs', 'the pinned map wins over the unknown tag')
    assert.equal(attribution.ownerOf('nope'), undefined, 'names outside the boot snapshot stay untagged')
  })

  test('a later runtime registration overrides the boot-time unknown tag', async () => {
    const app = new Context()
    const fake = {
      register() { return () => {} },
      layers: { global: { tools: { entries: () => [['claim_files', {}]] } } },
    }
    app.provide('tools', fake)
    const attribution = createToolAttribution(app)
    await app.plugin({
      name: 'claim-like',
      apply(sub) {
        ;((sub as any).tools as { register(def?: unknown): () => void }).register({ name: 'claim_files' })
      },
    })
    assert.equal(attribution.ownerOf('claim_files'), 'claim-like', 'a later live record beats the boot unknown tag')
  })

  test('tolerates tools values without an enumerable boot layer', () => {
    for (const tools of [
      undefined,
      { register() { return () => {} } },
      { register() { return () => {} }, layers: {} },
      { register() { return () => {} }, layers: { global: {} } },
      { register() { return () => {} }, layers: { global: { tools: {} } } },
      // A host tool service whose enumeration throws must not take the
      // caller's startup down — the boot snapshot just stays empty.
      {
        register() { return () => {} },
        layers: { global: { tools: { entries() { throw new Error('boom') } } } },
      },
    ]) {
      const app = new Context()
      if (tools !== undefined) app.provide('tools', tools)
      const attribution = createToolAttribution(app)
      assert.equal(attribution.ownerOf('whatever'), undefined, 'no boot entries → no unknown tag')
    }
  })

  test('resolves root-named local-link plugins through the register call stack', async () => {
    const app = new Context()
    const attribution = createToolAttribution(app)
    const fake = { register() { return () => {} } }
    app.provide('tools', fake)
    const { registerTool } = await import('./fixtures/local-plugin/index.js')
    await app.plugin({
      apply(sub) {
        registerTool((sub as any).tools)
      },
    })
    assert.equal(attribution.ownerOf('local_tool'), 'local-plugin')
  })

  test('callerPackageFrom resolves the registering package from a stack', () => {
    const anon = path.normalize(fileURLToPath(new URL('./fixtures/anon-plugin/index.js', import.meta.url)))
    const local = path.normalize(fileURLToPath(new URL('./fixtures/local-plugin/index.js', import.meta.url)))
    assert.equal(callerPackageFrom(undefined), undefined)
    assert.equal(callerPackageFrom('no frames here'), undefined)
    // non-frame lines are skipped
    assert.equal(callerPackageFrom(`Error\n  note\n  at ${local}:1:2`), 'local-plugin')
    // frames without a :line:col position are skipped, unresolvable targets yield undefined
    assert.equal(callerPackageFrom('Error\n    at x (file:///D:/no-position.js)\n    at y (file:///D:/has-position.js:1:2)'), undefined)
    // bare absolute path frames (transpiled modules) resolve
    assert.equal(callerPackageFrom(`Error\n    at ${local}:1:2\n    at ${specFile}:3:4`), 'local-plugin')
    // file:// URL frames with fn (...) wrappers resolve
    assert.equal(callerPackageFrom(`Error\n    at registerBlind (${fileUrl(anon)}:1:2)\n    at run (${fileUrl(specFile)}:3:4)`), 'anon-plugin')
    // frames resolving to this package are skipped, the next package wins
    assert.equal(callerPackageFrom(`Error\n    at ${specFile}:1:2\n    at ${local}:3:4`), 'local-plugin')
    // this module's own frames are skipped too
    const attributionFile = path.normalize(fileURLToPath(new URL('../../src/host/attribution.ts', import.meta.url)))
    assert.equal(callerPackageFrom(`Error\n    at ${attributionFile}:1:2\n    at ${local}:3:4`), 'local-plugin')
    // unparseable file URLs are skipped
    assert.equal(callerPackageFrom(`Error\n    at x (file:///D:/bad%zz.js:1:2)\n    at y (${fileUrl(local)}:3:4)`), 'local-plugin')
    // non-absolute targets (e.g. node: internal frames) are skipped
    assert.equal(callerPackageFrom(`Error\n    at ModuleJob.run (node:internal/modules/esm/module_job.js:346:20)\n    at y (${fileUrl(local)}:3:4)`), 'local-plugin')
    // async wrappers are tolerated
    assert.equal(callerPackageFrom(`Error\n    at async fn (${fileUrl(anon)}:1:2)`), 'anon-plugin')
  })

  test('packageNameFrom walks up to the nearest readable package.json', () => {
    const local = fileURLToPath(new URL('./fixtures/local-plugin/index.js', import.meta.url))
    assert.equal(packageNameFrom(local), 'local-plugin')
    assert.equal(packageNameFrom(local), 'local-plugin', 'results are cached per directory')
    const malformed = fileURLToPath(new URL('./fixtures/bad-pkg/index.js', import.meta.url))
    assert.equal(packageNameFrom(malformed), 'dsh-context', 'a malformed package.json is skipped, the walk continues upward')
    const rootless = path.join(path.parse(process.cwd()).root, '__dsh_context_nonexistent__', 'file.js')
    assert.equal(packageNameFrom(rootless), undefined, 'no readable boundary up to the filesystem root')
    const deep = path.join(path.parse(process.cwd()).root, ...Array.from({ length: 13 }, (_, i) => `__dsh_ctx_l${i}__`), 'file.js')
    assert.equal(packageNameFrom(deep), undefined, 'the walk-depth cap returns undefined before reaching the root')
  })
})