diff --git a/apps/api/src/federation/command.service.spec.ts b/apps/api/src/federation/command.service.spec.ts index 42e5dd7..a6ea75f 100644 --- a/apps/api/src/federation/command.service.spec.ts +++ b/apps/api/src/federation/command.service.spec.ts @@ -154,6 +154,15 @@ describe("CommandService", () => { signature: "signature-123", }) ); + + // Verify status was checked in the query + expect(prisma.federationConnection.findUnique).toHaveBeenCalledWith({ + where: { + id: mockConnectionId, + workspaceId: mockWorkspaceId, + status: FederationConnectionStatus.ACTIVE, + }, + }); }); it("should throw error if connection not found", async () => { @@ -165,19 +174,21 @@ describe("CommandService", () => { }); it("should throw error if connection is not active", async () => { - const mockConnection = { - id: mockConnectionId, - workspaceId: mockWorkspaceId, - status: FederationConnectionStatus.SUSPENDED, - }; - - vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue( - mockConnection as never - ); + // Connection should not be found by query because it's not ACTIVE + vi.spyOn(prisma.federationConnection, "findUnique").mockResolvedValue(null); await expect( service.sendCommand(mockWorkspaceId, mockConnectionId, "test", {}) - ).rejects.toThrow("Connection is not active"); + ).rejects.toThrow("Connection not found"); + + // Verify status was checked in the query + expect(prisma.federationConnection.findUnique).toHaveBeenCalledWith({ + where: { + id: mockConnectionId, + workspaceId: mockWorkspaceId, + status: FederationConnectionStatus.ACTIVE, + }, + }); }); it("should mark command as failed if sending fails", async () => { diff --git a/apps/api/src/federation/command.service.ts b/apps/api/src/federation/command.service.ts index 626c1af..e094a59 100644 --- a/apps/api/src/federation/command.service.ts +++ b/apps/api/src/federation/command.service.ts @@ -41,19 +41,19 @@ export class CommandService { commandType: string, payload: Record ): Promise { - // Validate connection exists and is active + // Validate connection exists and is active (enforced in query) const connection = await this.prisma.federationConnection.findUnique({ - where: { id: connectionId, workspaceId }, + where: { + id: connectionId, + workspaceId, + status: FederationConnectionStatus.ACTIVE, + }, }); if (!connection) { throw new Error("Connection not found"); } - if (connection.status !== FederationConnectionStatus.ACTIVE) { - throw new Error("Connection is not active"); - } - // Get local instance identity const identity = await this.federationService.getInstanceIdentity(); @@ -132,7 +132,7 @@ export class CommandService { throw new Error("Command timestamp is outside acceptable range"); } - // Find connection for remote instance + // Find connection for remote instance (status enforced in query) const connection = await this.prisma.federationConnection.findFirst({ where: { remoteInstanceId: commandMessage.instanceId, @@ -144,11 +144,6 @@ export class CommandService { throw new Error("No connection found for remote instance"); } - // Validate connection is active - if (connection.status !== FederationConnectionStatus.ACTIVE) { - throw new Error("Connection is not active"); - } - // Verify signature const { signature, ...messageToVerify } = commandMessage; const verificationResult = await this.signatureService.verifyMessage( diff --git a/apps/api/src/federation/query.service.spec.ts b/apps/api/src/federation/query.service.spec.ts index 8b1b59f..6d8b431 100644 --- a/apps/api/src/federation/query.service.spec.ts +++ b/apps/api/src/federation/query.service.spec.ts @@ -143,7 +143,11 @@ describe("QueryService", () => { expect(result.messageType).toBe(FederationMessageType.QUERY); expect(result.query).toBe(query); expect(mockPrisma.federationConnection.findUnique).toHaveBeenCalledWith({ - where: { id: connectionId, workspaceId }, + where: { + id: connectionId, + workspaceId, + status: FederationConnectionStatus.ACTIVE, + }, }); expect(mockPrisma.federationMessage.create).toHaveBeenCalled(); expect(mockHttpService.post).toHaveBeenCalledWith( @@ -168,17 +172,21 @@ describe("QueryService", () => { }); it("should throw error if connection not active", async () => { - const mockConnection = { - id: "connection-1", - workspaceId: "workspace-1", - status: FederationConnectionStatus.PENDING, - }; - - mockPrisma.federationConnection.findUnique.mockResolvedValue(mockConnection); + // Connection should not be found by query because it's not ACTIVE + mockPrisma.federationConnection.findUnique.mockResolvedValue(null); await expect( service.sendQuery("workspace-1", "connection-1", "SELECT * FROM tasks") - ).rejects.toThrow("Connection is not active"); + ).rejects.toThrow("Connection not found"); + + // Verify that findUnique was called with status: ACTIVE in the query + expect(mockPrisma.federationConnection.findUnique).toHaveBeenCalledWith({ + where: { + id: "connection-1", + workspaceId: "workspace-1", + status: FederationConnectionStatus.ACTIVE, + }, + }); }); it("should handle network errors gracefully", async () => { @@ -305,19 +313,21 @@ describe("QueryService", () => { signature: "valid-signature", }; - const mockConnection = { - id: "connection-1", - workspaceId: "workspace-1", - remoteInstanceId: "remote-instance-1", - status: FederationConnectionStatus.PENDING, - }; - - mockPrisma.federationConnection.findFirst.mockResolvedValue(mockConnection); + // Connection should not be found because status filter in query excludes non-ACTIVE + mockPrisma.federationConnection.findFirst.mockResolvedValue(null); mockSignatureService.validateTimestamp.mockReturnValue(true); await expect(service.handleIncomingQuery(queryMessage)).rejects.toThrow( - "Connection is not active" + "No connection found for remote instance" ); + + // Verify the findFirst was called with status: ACTIVE filter + expect(mockPrisma.federationConnection.findFirst).toHaveBeenCalledWith({ + where: { + remoteInstanceId: "remote-instance-1", + status: FederationConnectionStatus.ACTIVE, + }, + }); }); it("should reject query from unknown instance", async () => { diff --git a/apps/api/src/federation/query.service.ts b/apps/api/src/federation/query.service.ts index 6a2458b..a56c9e0 100644 --- a/apps/api/src/federation/query.service.ts +++ b/apps/api/src/federation/query.service.ts @@ -38,19 +38,19 @@ export class QueryService { query: string, context?: Record ): Promise { - // Validate connection exists and is active + // Validate connection exists and is active (enforced in query) const connection = await this.prisma.federationConnection.findUnique({ - where: { id: connectionId, workspaceId }, + where: { + id: connectionId, + workspaceId, + status: FederationConnectionStatus.ACTIVE, + }, }); if (!connection) { throw new Error("Connection not found"); } - if (connection.status !== FederationConnectionStatus.ACTIVE) { - throw new Error("Connection is not active"); - } - // Get local instance identity const identity = await this.federationService.getInstanceIdentity(); @@ -129,7 +129,7 @@ export class QueryService { throw new Error("Query timestamp is outside acceptable range"); } - // Find connection for remote instance + // Find connection for remote instance (status enforced in query) const connection = await this.prisma.federationConnection.findFirst({ where: { remoteInstanceId: queryMessage.instanceId, @@ -141,11 +141,6 @@ export class QueryService { throw new Error("No connection found for remote instance"); } - // Validate connection is active - if (connection.status !== FederationConnectionStatus.ACTIVE) { - throw new Error("Connection is not active"); - } - // Verify signature const { signature, ...messageToVerify } = queryMessage; const verificationResult = await this.signatureService.verifyMessage(