From 9642cd41d4afd4f01cfdf88b4b9d92637ad1c52c Mon Sep 17 00:00:00 2001 From: Jason Woltje Date: Sat, 7 Mar 2026 11:53:39 -0600 Subject: [PATCH] feat(orchestrator): add subagent session tree endpoint --- .../agent-ingestion.service.ts | 4 +-- .../src/api/agents/agent-tree.service.ts | 30 +++++++++++++++++ .../agents-killswitch.controller.spec.ts | 11 ++++++- .../src/api/agents/agents.controller.spec.ts | 32 ++++++++++++++++++- .../src/api/agents/agents.controller.ts | 13 +++++++- .../src/api/agents/agents.module.ts | 2 ++ .../api/agents/dto/agent-tree-response.dto.ts | 9 ++++++ .../src/api/agents/dto/spawn-agent.dto.ts | 4 +++ .../src/spawner/agent-spawner.service.ts | 8 ++++- .../src/spawner/types/agent-spawner.types.ts | 2 ++ 10 files changed, 109 insertions(+), 6 deletions(-) create mode 100644 apps/orchestrator/src/api/agents/agent-tree.service.ts create mode 100644 apps/orchestrator/src/api/agents/dto/agent-tree-response.dto.ts diff --git a/apps/orchestrator/src/agent-ingestion/agent-ingestion.service.ts b/apps/orchestrator/src/agent-ingestion/agent-ingestion.service.ts index 64412ab..0e0b00a 100644 --- a/apps/orchestrator/src/agent-ingestion/agent-ingestion.service.ts +++ b/apps/orchestrator/src/agent-ingestion/agent-ingestion.service.ts @@ -25,14 +25,14 @@ export class AgentIngestionService { where: { sessionId: agentId }, create: { sessionId: agentId, - parentSessionId: parentAgentId, + parentSessionId: parentAgentId ?? null, missionId, taskId, agentType, status: "spawning", }, update: { - parentSessionId: parentAgentId, + parentSessionId: parentAgentId ?? null, missionId, taskId, agentType, diff --git a/apps/orchestrator/src/api/agents/agent-tree.service.ts b/apps/orchestrator/src/api/agents/agent-tree.service.ts new file mode 100644 index 0000000..6a5ed68 --- /dev/null +++ b/apps/orchestrator/src/api/agents/agent-tree.service.ts @@ -0,0 +1,30 @@ +import { Injectable } from "@nestjs/common"; +import { PrismaService } from "../../prisma/prisma.service"; +import { AgentTreeResponseDto } from "./dto/agent-tree-response.dto"; + +@Injectable() +export class AgentTreeService { + constructor(private readonly prisma: PrismaService) {} + + async getTree(): Promise { + const entries = await this.prisma.agentSessionTree.findMany({ + orderBy: { spawnedAt: "desc" }, + take: 200, + }); + + const response: AgentTreeResponseDto[] = []; + for (const entry of entries) { + response.push({ + sessionId: entry.sessionId, + parentSessionId: entry.parentSessionId ?? null, + status: entry.status, + agentType: entry.agentType ?? null, + taskSource: entry.taskSource ?? null, + spawnedAt: entry.spawnedAt.toISOString(), + completedAt: entry.completedAt?.toISOString() ?? null, + }); + } + + return response; + } +} diff --git a/apps/orchestrator/src/api/agents/agents-killswitch.controller.spec.ts b/apps/orchestrator/src/api/agents/agents-killswitch.controller.spec.ts index 59abe78..a3368b1 100644 --- a/apps/orchestrator/src/api/agents/agents-killswitch.controller.spec.ts +++ b/apps/orchestrator/src/api/agents/agents-killswitch.controller.spec.ts @@ -7,6 +7,7 @@ import { KillswitchService } from "../../killswitch/killswitch.service"; import { AgentEventsService } from "./agent-events.service"; import { AgentMessagesService } from "./agent-messages.service"; import { AgentControlService } from "./agent-control.service"; +import { AgentTreeService } from "./agent-tree.service"; import type { KillAllResult } from "../../killswitch/killswitch.service"; describe("AgentsController - Killswitch Endpoints", () => { @@ -41,6 +42,9 @@ describe("AgentsController - Killswitch Endpoints", () => { pauseAgent: ReturnType; resumeAgent: ReturnType; }; + let mockTreeService: { + getTree: ReturnType; + }; beforeEach(() => { mockKillswitchService = { @@ -89,6 +93,10 @@ describe("AgentsController - Killswitch Endpoints", () => { resumeAgent: vi.fn().mockResolvedValue(undefined), }; + mockTreeService = { + getTree: vi.fn().mockResolvedValue([]), + }; + controller = new AgentsController( mockQueueService as unknown as QueueService, mockSpawnerService as unknown as AgentSpawnerService, @@ -96,7 +104,8 @@ describe("AgentsController - Killswitch Endpoints", () => { mockKillswitchService as unknown as KillswitchService, mockEventsService as unknown as AgentEventsService, mockMessagesService as unknown as AgentMessagesService, - mockControlService as unknown as AgentControlService + mockControlService as unknown as AgentControlService, + mockTreeService as unknown as AgentTreeService ); }); diff --git a/apps/orchestrator/src/api/agents/agents.controller.spec.ts b/apps/orchestrator/src/api/agents/agents.controller.spec.ts index 5b30e86..f7ba182 100644 --- a/apps/orchestrator/src/api/agents/agents.controller.spec.ts +++ b/apps/orchestrator/src/api/agents/agents.controller.spec.ts @@ -6,6 +6,7 @@ import { KillswitchService } from "../../killswitch/killswitch.service"; import { AgentEventsService } from "./agent-events.service"; import { AgentMessagesService } from "./agent-messages.service"; import { AgentControlService } from "./agent-control.service"; +import { AgentTreeService } from "./agent-tree.service"; import { describe, it, expect, beforeEach, afterEach, vi } from "vitest"; describe("AgentsController", () => { @@ -42,6 +43,9 @@ describe("AgentsController", () => { pauseAgent: ReturnType; resumeAgent: ReturnType; }; + let treeService: { + getTree: ReturnType; + }; beforeEach(() => { // Create mock services @@ -93,6 +97,10 @@ describe("AgentsController", () => { resumeAgent: vi.fn().mockResolvedValue(undefined), }; + treeService = { + getTree: vi.fn().mockResolvedValue([]), + }; + // Create controller with mocked services controller = new AgentsController( queueService as unknown as QueueService, @@ -101,7 +109,8 @@ describe("AgentsController", () => { killswitchService as unknown as KillswitchService, eventsService as unknown as AgentEventsService, messagesService as unknown as AgentMessagesService, - controlService as unknown as AgentControlService + controlService as unknown as AgentControlService, + treeService as unknown as AgentTreeService ); }); @@ -113,6 +122,27 @@ describe("AgentsController", () => { expect(controller).toBeDefined(); }); + describe("getAgentTree", () => { + it("should return tree entries", async () => { + const entries = [ + { + sessionId: "agent-1", + parentSessionId: null, + status: "running", + agentType: "worker", + taskSource: "internal", + spawnedAt: "2026-03-07T00:00:00.000Z", + completedAt: null, + }, + ]; + + treeService.getTree.mockResolvedValue(entries); + + await expect(controller.getAgentTree()).resolves.toEqual(entries); + expect(treeService.getTree).toHaveBeenCalledTimes(1); + }); + }); + describe("listAgents", () => { it("should return empty array when no agents exist", () => { // Arrange diff --git a/apps/orchestrator/src/api/agents/agents.controller.ts b/apps/orchestrator/src/api/agents/agents.controller.ts index aa82360..f6304b9 100644 --- a/apps/orchestrator/src/api/agents/agents.controller.ts +++ b/apps/orchestrator/src/api/agents/agents.controller.ts @@ -30,6 +30,8 @@ import { AgentEventsService } from "./agent-events.service"; import { GetMessagesQueryDto } from "./dto/get-messages-query.dto"; import { AgentMessagesService } from "./agent-messages.service"; import { AgentControlService } from "./agent-control.service"; +import { AgentTreeService } from "./agent-tree.service"; +import { AgentTreeResponseDto } from "./dto/agent-tree-response.dto"; import { InjectAgentDto } from "./dto/inject-agent.dto"; import { PauseAgentDto, ResumeAgentDto } from "./dto/control-agent.dto"; @@ -56,7 +58,8 @@ export class AgentsController { private readonly killswitchService: KillswitchService, private readonly eventsService: AgentEventsService, private readonly messagesService: AgentMessagesService, - private readonly agentControlService: AgentControlService + private readonly agentControlService: AgentControlService, + private readonly agentTreeService: AgentTreeService ) {} /** @@ -78,6 +81,7 @@ export class AgentsController { // Spawn agent using spawner service const spawnResponse = this.spawnerService.spawnAgent({ taskId: dto.taskId, + ...(dto.parentAgentId !== undefined ? { parentAgentId: dto.parentAgentId } : {}), agentType: dto.agentType, context: dto.context, }); @@ -152,6 +156,13 @@ export class AgentsController { }; } + @Get("tree") + @UseGuards(OrchestratorApiKeyGuard) + @Throttle({ default: { limit: 200, ttl: 60000 } }) + async getAgentTree(): Promise { + return this.agentTreeService.getTree(); + } + /** * List all agents * @returns Array of all agent sessions with their status diff --git a/apps/orchestrator/src/api/agents/agents.module.ts b/apps/orchestrator/src/api/agents/agents.module.ts index 903f226..a1d5d76 100644 --- a/apps/orchestrator/src/api/agents/agents.module.ts +++ b/apps/orchestrator/src/api/agents/agents.module.ts @@ -9,6 +9,7 @@ import { AgentEventsService } from "./agent-events.service"; import { PrismaModule } from "../../prisma/prisma.module"; import { AgentMessagesService } from "./agent-messages.service"; import { AgentControlService } from "./agent-control.service"; +import { AgentTreeService } from "./agent-tree.service"; @Module({ imports: [QueueModule, SpawnerModule, KillswitchModule, ValkeyModule, PrismaModule], @@ -18,6 +19,7 @@ import { AgentControlService } from "./agent-control.service"; AgentEventsService, AgentMessagesService, AgentControlService, + AgentTreeService, ], }) export class AgentsModule {} diff --git a/apps/orchestrator/src/api/agents/dto/agent-tree-response.dto.ts b/apps/orchestrator/src/api/agents/dto/agent-tree-response.dto.ts new file mode 100644 index 0000000..25abf05 --- /dev/null +++ b/apps/orchestrator/src/api/agents/dto/agent-tree-response.dto.ts @@ -0,0 +1,9 @@ +export class AgentTreeResponseDto { + sessionId!: string; + parentSessionId!: string | null; + status!: string; + agentType!: string | null; + taskSource!: string | null; + spawnedAt!: string; + completedAt!: string | null; +} diff --git a/apps/orchestrator/src/api/agents/dto/spawn-agent.dto.ts b/apps/orchestrator/src/api/agents/dto/spawn-agent.dto.ts index 0bcd13b..1797f0b 100644 --- a/apps/orchestrator/src/api/agents/dto/spawn-agent.dto.ts +++ b/apps/orchestrator/src/api/agents/dto/spawn-agent.dto.ts @@ -116,6 +116,10 @@ export class SpawnAgentDto { @IsOptional() @IsIn(["strict", "standard", "minimal", "custom"]) gateProfile?: GateProfileType; + + @IsOptional() + @IsString() + parentAgentId?: string; } /** diff --git a/apps/orchestrator/src/spawner/agent-spawner.service.ts b/apps/orchestrator/src/spawner/agent-spawner.service.ts index 9c9ebba..b9f0805 100644 --- a/apps/orchestrator/src/spawner/agent-spawner.service.ts +++ b/apps/orchestrator/src/spawner/agent-spawner.service.ts @@ -115,7 +115,13 @@ export class AgentSpawnerService implements OnModuleDestroy { } void this.agentIngestionService - .recordAgentSpawned(agentId, undefined, undefined, request.taskId, request.agentType) + .recordAgentSpawned( + agentId, + request.parentAgentId, + undefined, + request.taskId, + request.agentType + ) .catch((error: unknown) => { const errorMessage = error instanceof Error ? error.message : String(error); this.logger.error(`Failed to record spawned ingestion for ${agentId}: ${errorMessage}`); diff --git a/apps/orchestrator/src/spawner/types/agent-spawner.types.ts b/apps/orchestrator/src/spawner/types/agent-spawner.types.ts index 079c9e6..2d3fba4 100644 --- a/apps/orchestrator/src/spawner/types/agent-spawner.types.ts +++ b/apps/orchestrator/src/spawner/types/agent-spawner.types.ts @@ -40,6 +40,8 @@ export interface SpawnAgentOptions { export interface SpawnAgentRequest { /** Unique task identifier */ taskId: string; + /** Optional parent session identifier for subagent lineage */ + parentAgentId?: string; /** Type of agent to spawn */ agentType: AgentType; /** Context for task execution */