Fix image removal lookup

Sources/Services/ContainerAPIService/Client/ClientImage.swift

@@ -170,7 +170,7 @@ extension ClientImage {
         var found: [ClientImage] = []
         for name in names {
             do {
-                guard let img = try Self._search(reference: name, in: all, containerSystemConfig: containerSystemConfig) else {
+                guard let img = try Self.match(reference: name, in: all, containerSystemConfig: containerSystemConfig) else {
                     errors.append(name)
                     continue
                 }
@@ -184,7 +184,7 @@ extension ClientImage {
 
     public static func get(reference: String, containerSystemConfig: ContainerSystemConfig) async throws -> ClientImage {
         let all = try await self.list()
-        guard let found = try self._search(reference: reference, in: all, containerSystemConfig: containerSystemConfig) else {
+        guard let found = try self.match(reference: reference, in: all, containerSystemConfig: containerSystemConfig) else {
             throw ContainerizationError(.notFound, message: "image with reference \(reference)")
         }
         return found
@@ -217,7 +217,15 @@ extension ClientImage {
         return 0
     }
 
-    private static func _search(reference: String, in all: [ClientImage], containerSystemConfig: ContainerSystemConfig) throws -> ClientImage? {
+    static func match(reference: String, in all: [ClientImage], containerSystemConfig: ContainerSystemConfig) throws -> ClientImage? {
+        if let image = all.first(where: { $0.reference == reference }) {
+            return image
+        }
+
+        if let image = Self.matchDigestPrefix(reference, in: all) {
+            return image
+        }
+
         let locallyBuiltImage = try {
             // Check if we have an image whose index descriptor contains the image name
             // as an annotation. Prefer this in all cases, since these are locally built images.
@@ -235,13 +243,78 @@ extension ClientImage {
         if let locallyBuiltImage {
             return locallyBuiltImage
         }
-        // If we don't find a match, try matching `ImageDescription.name` against the given
-        // input string, while also checking against its normalized form.
-        // Return the first match.
-        let normalizedReference = try Self.normalizeReference(reference, containerSystemConfig: containerSystemConfig)
-        return all.first(where: { image in
-            image.reference == reference || image.reference == normalizedReference
-        })
+
+        if let normalizedReference = try? Self.normalizeReference(reference, containerSystemConfig: containerSystemConfig),
+            let image = all.first(where: { $0.reference == normalizedReference })
+        {
+            return image
+        }
+
+        return try Self.matchDisplayName(reference, in: all, containerSystemConfig: containerSystemConfig)
+    }
+
+    private static func matchDigestPrefix(_ reference: String, in all: [ClientImage]) -> ClientImage? {
+        guard let digestPrefix = Self.digestPrefix(from: reference) else {
+            return nil
+        }
+        return Self.uniqueMatch(in: all) { image in
+            Self.digestIdentifier(from: image.digest).hasPrefix(digestPrefix)
+        }
+    }
+
+    private static func digestPrefix(from reference: String) -> String? {
+        let prefix: String
+        if reference.hasPrefix("sha256:") {
+            prefix = String(reference.dropFirst("sha256:".count))
+        } else {
+            guard !reference.contains(":") && !reference.contains("/") && !reference.contains("@") else {
+                return nil
+            }
+            prefix = reference
+        }
+        guard prefix.count >= 12, prefix.allSatisfy(\.isHexDigit) else {
+            return nil
+        }
+        return prefix.lowercased()
+    }
+
+    private static func digestIdentifier(from digest: String) -> String {
+        guard let separator = digest.firstIndex(of: ":") else {
+            return digest.lowercased()
+        }
+        return digest[digest.index(after: separator)...].lowercased()
+    }
+
+    private static func matchDisplayName(_ reference: String, in all: [ClientImage], containerSystemConfig: ContainerSystemConfig) throws -> ClientImage? {
+        let requestedReference = try Reference.parse(reference)
+        guard requestedReference.tag == nil && requestedReference.digest == nil else {
+            return nil
+        }
+
+        return try Self.uniqueMatch(in: all) { image in
+            let storedReference = try Reference.parse(image.reference)
+            if storedReference.name == requestedReference.name {
+                return true
+            }
+
+            let displayReference = try Self.denormalizeReference(image.reference, containerSystemConfig: containerSystemConfig)
+            let parsedDisplayReference = try Reference.parse(displayReference)
+            return parsedDisplayReference.name == requestedReference.name
+        }
+    }
+
+    private static func uniqueMatch(in images: [ClientImage], where predicate: (ClientImage) throws -> Bool) rethrows -> ClientImage? {
+        var match: ClientImage?
+        for image in images {
+            guard try predicate(image) else {
+                continue
+            }
+            if match != nil {
+                return nil
+            }
+            match = image
+        }
+        return match
     }
 
     public static func pull(

Tests/ContainerAPIClientTests/ClientImageLookupTests.swift

@@ -0,0 +1,137 @@
+//===----------------------------------------------------------------------===//
+// Copyright © 2026 Apple Inc. and the container project authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//   https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//===----------------------------------------------------------------------===//
+
+import ContainerPersistence
+import ContainerResource
+import ContainerizationOCI
+import Testing
+
+@testable import ContainerAPIClient
+
+struct ClientImageLookupTests {
+    private let containerSystemConfig = ContainerSystemConfig()
+
+    @Test
+    func testMatchUsesDisplayedDefaultRegistryNameWhenUnique() throws {
+        let ubuntu = Self.image(
+            reference: "docker.io/library/ubuntu:22.04",
+            digest: "01a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+        let fedora = Self.image(
+            reference: "docker.io/library/fedora:latest",
+            digest: "11a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "ubuntu", in: [fedora, ubuntu], containerSystemConfig: containerSystemConfig)
+
+        #expect(match?.reference == ubuntu.reference)
+    }
+
+    @Test
+    func testMatchDoesNotGuessAmbiguousDisplayedNames() throws {
+        let jammy = Self.image(
+            reference: "docker.io/library/ubuntu:22.04",
+            digest: "21a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+        let noble = Self.image(
+            reference: "docker.io/library/ubuntu:24.04",
+            digest: "31a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "ubuntu", in: [jammy, noble], containerSystemConfig: containerSystemConfig)
+
+        #expect(match == nil)
+    }
+
+    @Test
+    func testMatchUsesDigestPrefixFromImageList() throws {
+        let ubuntu = Self.image(
+            reference: "docker.io/library/ubuntu:22.04",
+            digest: "e10577f0db681cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "e10577f0db68", in: [ubuntu], containerSystemConfig: containerSystemConfig)
+
+        #expect(match?.reference == ubuntu.reference)
+    }
+
+    @Test
+    func testMatchDoesNotTreatHexTagsAsDigestPrefixes() throws {
+        let digestMatch = Self.image(
+            reference: "docker.io/library/fedora:latest",
+            digest: "deadbeefdead1cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+        let tagged = Self.image(
+            reference: "docker.io/library/foo:deadbeefdead",
+            digest: "71a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "foo:deadbeefdead", in: [digestMatch, tagged], containerSystemConfig: containerSystemConfig)
+
+        #expect(match?.reference == tagged.reference)
+    }
+
+    @Test
+    func testMatchDoesNotGuessAmbiguousDigestPrefixes() throws {
+        let first = Self.image(
+            reference: "docker.io/library/ubuntu:22.04",
+            digest: "d56a2534ffd21cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+        let second = Self.image(
+            reference: "docker.io/library/fedora:latest",
+            digest: "d56a2534ffd22cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "d56a2534ffd2", in: [first, second], containerSystemConfig: containerSystemConfig)
+
+        #expect(match == nil)
+    }
+
+    @Test
+    func testMatchKeepsNormalizedReferenceBehavior() throws {
+        let ubuntu = Self.image(
+            reference: "docker.io/library/ubuntu:latest",
+            digest: "41a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+
+        let match = try ClientImage.match(reference: "ubuntu", in: [ubuntu], containerSystemConfig: containerSystemConfig)
+
+        #expect(match?.reference == ubuntu.reference)
+    }
+
+    @Test
+    func testMatchPreservesLocalImageAnnotationPreference() throws {
+        let remote = Self.image(
+            reference: "docker.io/library/foo:latest",
+            digest: "51a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000"
+        )
+        let local = Self.image(
+            reference: "registry.local/builds/foo:latest",
+            digest: "61a3ee0b5e413cefaaffc6ab0000000000000000000000000000000000000000",
+            annotations: [AnnotationKeys.containerizationImageName: "foo:latest"]
+        )
+
+        let match = try ClientImage.match(reference: "foo", in: [remote, local], containerSystemConfig: containerSystemConfig)
+
+        #expect(match?.reference == local.reference)
+    }
+
+    private static func image(reference: String, digest: String, annotations: [String: String]? = nil) -> ClientImage {
+        let descriptor = Descriptor(mediaType: MediaTypes.index, digest: "sha256:\(digest)", size: 0, annotations: annotations)
+        let description = ImageDescription(reference: reference, descriptor: descriptor)
+        return ClientImage(description: description)
+    }
+}