stack/apps/gateway/src/agent/__tests__/agent-service-ownership.test.ts

import { ForbiddenException } from '@nestjs/common';
import { describe, expect, it, vi } from 'vitest';
import { AgentService, type AgentSession } from '../agent.service.js';
import type { ActorTenantScope } from '../../auth/session-scope.js';

const CONVERSATION_ID = '22222222-2222-4222-8222-222222222222';
const OWNER_SCOPE: ActorTenantScope = { userId: 'owner-user', tenantId: 'owner-tenant' };
const FOREIGN_SCOPE: ActorTenantScope = { userId: 'foreign-user', tenantId: 'foreign-tenant' };

type AgentServiceInternals = {
  sessions: Map<string, AgentSession>;
  creating: Map<string, Promise<AgentSession>>;
};

function makeService(operatorMemory: unknown = null): AgentService {
  return new AgentService(
    {
      getDefaultModel: vi.fn(() => null),
      getRegistry: vi.fn(() => ({})),
      findModel: vi.fn(),
      listAvailableModels: vi.fn(() => []),
    } as never,
    {} as never,
    {} as never,
    { available: false } as never,
    {} as never,
    { getToolDefinitions: vi.fn(() => []) } as never,
    { loadForSession: vi.fn(async () => ({ metaTools: [], promptAdditions: [] })) } as never,
    null,
    null,
    { collect: vi.fn().mockResolvedValue(undefined) } as never,
    operatorMemory as never,
  );
}

function internals(service: AgentService): AgentServiceInternals {
  return service as unknown as AgentServiceInternals;
}

function makeSession(scope: ActorTenantScope = OWNER_SCOPE): AgentSession {
  return {
    id: CONVERSATION_ID,
    provider: 'test-provider',
    modelId: 'test-model',
    piSession: {
      thinkingLevel: 'off',
      getAvailableThinkingLevels: vi.fn().mockReturnValue(['off', 'low', 'high']),
      setThinkingLevel: vi.fn(),
      abort: vi.fn().mockResolvedValue(undefined),
      prompt: vi.fn().mockResolvedValue(undefined),
      dispose: vi.fn(),
      getSessionStats: vi.fn(),
      getContextUsage: vi.fn(),
    } as unknown as AgentSession['piSession'],
    listeners: new Set(),
    unsubscribe: vi.fn(),
    createdAt: Date.now(),
    promptCount: 0,
    channels: new Set(),
    skillPromptAdditions: [],
    sandboxDir: '/tmp/tess-session-ownership-test',
    allowedTools: null,
    userId: scope.userId,
    tenantId: scope.tenantId,
    metrics: {
      tokens: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
      modelSwitches: 0,
      messageCount: 0,
      lastActivityAt: new Date('2026-07-12T00:00:00Z').toISOString(),
    },
  };
}

describe('AgentService owner/tenant scope enforcement', () => {
  it('allows owner-scoped operations and rejects foreign scopes for seeded sessions', async () => {
    const service = makeService();
    const session = makeSession();
    internals(service).sessions.set(CONVERSATION_ID, session);

    expect(service.getSession(CONVERSATION_ID, OWNER_SCOPE)).toBe(session);
    expect(service.getSession(CONVERSATION_ID, FOREIGN_SCOPE)).toBeUndefined();
    expect(service.getSessionInfo(CONVERSATION_ID, FOREIGN_SCOPE)).toBeUndefined();
    expect(service.listSessions(OWNER_SCOPE)).toHaveLength(1);
    expect(service.listSessions(FOREIGN_SCOPE)).toEqual([]);

    service.addChannel(CONVERSATION_ID, 'websocket:owner', OWNER_SCOPE);
    expect(session.channels.has('websocket:owner')).toBe(true);
    expect(() => service.addChannel(CONVERSATION_ID, 'websocket:foreign', FOREIGN_SCOPE)).toThrow(
      ForbiddenException,
    );
    expect(() => service.removeChannel(CONVERSATION_ID, 'websocket:owner', FOREIGN_SCOPE)).toThrow(
      ForbiddenException,
    );

    expect(() =>
      service.updateSessionModel(CONVERSATION_ID, 'foreign-model', FOREIGN_SCOPE),
    ).toThrow(ForbiddenException);
    service.updateSessionModel(CONVERSATION_ID, 'owner-model', OWNER_SCOPE);
    expect(session.modelId).toBe('owner-model');

    expect(() =>
      service.applyAgentConfig(CONVERSATION_ID, 'agent-foreign', 'Foreign Agent', FOREIGN_SCOPE),
    ).toThrow(ForbiddenException);
    service.applyAgentConfig(CONVERSATION_ID, 'agent-owner', 'Owner Agent', OWNER_SCOPE);
    expect(session.agentConfigId).toBe('agent-owner');

    expect(() => service.onEvent(CONVERSATION_ID, vi.fn(), FOREIGN_SCOPE)).toThrow(
      ForbiddenException,
    );
    const cleanup = service.onEvent(CONVERSATION_ID, vi.fn(), OWNER_SCOPE);
    cleanup();

    await expect(
      service.prompt(CONVERSATION_ID, 'foreign prompt', FOREIGN_SCOPE),
    ).rejects.toBeInstanceOf(ForbiddenException);
    await service.prompt(CONVERSATION_ID, 'owner prompt', OWNER_SCOPE);
    expect(session.piSession.prompt).toHaveBeenCalledWith('owner prompt');

    await expect(service.destroySession(CONVERSATION_ID, FOREIGN_SCOPE)).rejects.toBeInstanceOf(
      ForbiddenException,
    );
    expect(internals(service).sessions.has(CONVERSATION_ID)).toBe(true);

    await service.destroySession(CONVERSATION_ID, OWNER_SCOPE);
    expect(session.piSession.dispose).toHaveBeenCalled();
    expect(internals(service).sessions.has(CONVERSATION_ID)).toBe(false);
  });

  it('derives the operator-memory scope on the createSession production path', async () => {
    const plugin = { capture: vi.fn(), search: vi.fn() };
    const service = makeService(plugin);
    const buildTools = vi.spyOn(service as never, 'buildToolsForSandbox').mockReturnValue([]);

    // Session construction reaches the real scope derivation before the intentionally incomplete
    // Pi test double rejects later in createAgentSession.
    await service.createSession(CONVERSATION_ID, OWNER_SCOPE).catch(() => undefined);

    expect(buildTools).toHaveBeenCalledWith(expect.any(String), OWNER_SCOPE.userId, {
      tenantId: OWNER_SCOPE.tenantId,
      ownerId: OWNER_SCOPE.userId,
      sessionId: CONVERSATION_ID,
    });
  });

  it('denies a foreign actor before it can obtain another session operator-memory scope', async () => {
    const plugin = { capture: vi.fn(), search: vi.fn() };
    const service = makeService(plugin);
    internals(service).sessions.set(CONVERSATION_ID, makeSession());
    const buildTools = vi.spyOn(service as never, 'buildToolsForSandbox');

    await expect(service.createSession(CONVERSATION_ID, FOREIGN_SCOPE)).rejects.toBeInstanceOf(
      ForbiddenException,
    );

    expect(buildTools).not.toHaveBeenCalled();
    expect(plugin.capture).not.toHaveBeenCalled();
    expect(plugin.search).not.toHaveBeenCalled();
  });

  it('checks owner/tenant scope before returning an in-flight session creation', async () => {
    const service = makeService();
    const session = makeSession();
    internals(service).creating.set(CONVERSATION_ID, Promise.resolve(session));

    await expect(
      service.createSession(CONVERSATION_ID, {
        userId: FOREIGN_SCOPE.userId,
        tenantId: FOREIGN_SCOPE.tenantId,
      }),
    ).rejects.toBeInstanceOf(ForbiddenException);

    await expect(
      service.createSession(CONVERSATION_ID, {
        userId: OWNER_SCOPE.userId,
        tenantId: OWNER_SCOPE.tenantId,
      }),
    ).resolves.toBe(session);
  });
});