feat(rendition): authorize bounded non-PDF Mistral OCR inputs

document/mistral/policy_authorization_test.go

@@ -0,0 +1,20 @@
+package mistral
+
+import "testing"
+
+func TestPolicyDoesNotAuthorizeUnprovedNonPDFFormats(t *testing.T) {
+	policy := testPolicy(t, 1<<20, 10)
+	manifest := syntheticManifest(t, policy, true)
+
+	for _, formatID := range []string{"docx", "pptx", "xlsx", "odt", "epub", "txt"} {
+		t.Run(formatID, func(t *testing.T) {
+			if expectedUnitBound(formatID) == UnitBoundLocalExact || localUnitCounters[formatID] != nil {
+				t.Fatalf("unproved format %q is registered for local authority", formatID)
+			}
+			_, err := policy.Authorize(manifest, formatID)
+			if err == nil {
+				t.Fatalf("Policy.Authorize(%q) succeeded without provider-authentic unit evidence", formatID)
+			}
+		})
+	}
+}

document/mistral/rendition.go

@@ -156,14 +156,18 @@ func (client *RenditionClient) Render(
 	if err != nil {
 		return document.RenditionResult{}, err
 	}
-	providerResult, err := client.process(operation, source, metadata, candidate,
+	providerResult, err := client.process(operation, source, metadata, candidate, localUnits,
 		min(client.policy.values.MaxResponseBytes, int64(authorization.MaxTotalResultBytes)))
 	if err != nil {
 		return document.RenditionResult{}, err
 	}
-	if candidate.ID == formatIDPDF && int64(providerResult.UnitsProcessed) != localUnits {
+	if (candidate.ID == formatIDPDF || localUnits > 0) && int64(providerResult.UnitsProcessed) != localUnits {
+		message := "Mistral OCR unit count changed"
+		if candidate.ID == formatIDPDF {
+			message = "Mistral OCR page count changed"
+		}
 		return document.RenditionResult{}, renditionProvider.Classified(document.RenditionErrorPolicyRejected,
-			"Mistral OCR page count changed", ErrCapabilityContract)
+			message, ErrCapabilityContract)
 	}
 	completedAt := time.Now().UTC()
 	if err := operation.Check(); err != nil {
@@ -199,8 +203,8 @@ func (client *RenditionClient) Render(
 	return document.RenditionResult{Evidence: evidence, ProviderMarkdown: markdown, Receipt: receipt}, nil
 }
 
-// verifySource re-detects the exact format and proves the PDF page count
-// before any byte leaves the process.
+// verifySource re-detects the exact format and proves any registered local
+// unit count before any byte leaves the process.
 func (client *RenditionClient) verifySource(
 	source []byte, metadata document.AuthorizedUploadMetadata,
 ) (CandidateFormat, int64, error) {
@@ -213,24 +217,36 @@ func (client *RenditionClient) verifySource(
 		return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorPolicyRejected,
 			"Mistral input identity does not match authorization", nil)
 	}
-	if candidate.ID != formatIDPDF {
+	if candidate.ID == formatIDPDF {
+		localUnits, err := formatdetect.CountPDFPages(source)
+		if err != nil {
+			return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorUnsupportedInput,
+				"Mistral PDF page count could not be verified", err)
+		}
+		if localUnits <= 0 || localUnits > int64(client.policy.values.MaxUnits) {
+			return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorPolicyRejected,
+				"Mistral PDF exceeds the complete unit limit", nil)
+		}
+		return candidate, localUnits, nil
+	}
+	if expectedUnitBound(candidate.ID) != UnitBoundLocalExact {
 		return candidate, 0, nil
 	}
-	localUnits, err := formatdetect.CountPDFPages(source)
+	localUnits, err := countLocalUnits(candidate, bytes.NewReader(source), int64(len(source)))
 	if err != nil {
 		return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorUnsupportedInput,
-			"Mistral PDF page count could not be verified", err)
+			"Mistral OCR local unit count could not be verified", err)
 	}
-	if localUnits <= 0 || localUnits > int64(client.policy.values.MaxUnits) {
+	if localUnits <= 0 || int64(localUnits) > int64(client.policy.values.MaxUnits) {
 		return CandidateFormat{}, 0, renditionProvider.Classified(document.RenditionErrorPolicyRejected,
-			"Mistral PDF exceeds the complete unit limit", nil)
+			"Mistral OCR exceeds the complete unit limit", nil)
 	}
-	return candidate, localUnits, nil
+	return candidate, int64(localUnits), nil
 }
 
 func (client *RenditionClient) process(
 	operation *providerutil.Operation, source []byte, metadata document.AuthorizedUploadMetadata,
-	candidate CandidateFormat, maxResponseBytes int64,
+	candidate CandidateFormat, localUnits int64, maxResponseBytes int64,
 ) (Result, error) {
 	formatAuthorization, err := client.policy.Authorize(client.manifest, candidate.ID)
 	if err != nil {
@@ -239,7 +255,7 @@ func (client *RenditionClient) process(
 	}
 	snapshot := preparedSnapshot{
 		size: int64(len(source)), sha256: metadata.SHA256, format: candidate,
-		mediaType: metadata.MediaType,
+		mediaType: metadata.MediaType, localUnits: int(localUnits),
 	}
 	snapshotForAttempt := func() (preparedSnapshot, error) {
 		digest := sha256.Sum256(source)

document/mistral/rendition_test.go

@@ -114,6 +114,77 @@ func TestRenditionClientMapsExactMistralOCRResponse(t *testing.T) {
 	assert.NotContains(t, fmt.Sprintf("%+v", result.Receipt), "synthetic-key")
 }
 
+func TestRenditionClientUsesLocalExactUnitsForAuthorizedNonPDF(t *testing.T) {
+	policy := testPolicy(t, 1<<20, 10)
+	manifest := syntheticManifest(t, policy, true)
+
+	previousMethod, hadMethod := expectedUnitBounds["docx"]
+	previousCounter, hadCounter := localUnitCounters["docx"]
+	t.Cleanup(func() {
+		if hadMethod {
+			expectedUnitBounds["docx"] = previousMethod
+		} else {
+			delete(expectedUnitBounds, "docx")
+		}
+		if hadCounter {
+			localUnitCounters["docx"] = previousCounter
+		} else {
+			delete(localUnitCounters, "docx")
+		}
+	})
+	expectedUnitBounds["docx"] = UnitBoundLocalExact
+	localUnitCounters["docx"] = func(io.ReaderAt, int64) (int, error) { return 2, nil }
+	for index := range manifest.Results {
+		if manifest.Results[index].FormatID == "docx" {
+			manifest.Results[index].UnitCount = 2
+			manifest.Results[index].UnitsProcessed = 2
+			manifest.Results[index].LocalUnits = 2
+			manifest.Results[index].UnitBoundMethod = UnitBoundLocalExact
+		}
+	}
+
+	descriptor := renditionDescriptor(t, policy, manifest, "docx")
+	response := fmt.Sprintf(
+		`{"model":"mistral-ocr-4-0","pages":[{"index":0,"markdown":"first"},{"index":1,"markdown":"second"}],"usage_info":{"pages_processed":2,"doc_size_bytes":%d}}`,
+		len(documentZIP(t, map[string]string{
+			ooxmlContentTypesName: docxContentTypes("application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"),
+			"word/document.xml":   "<document/>",
+		})),
+	)
+	source := documentZIP(t, map[string]string{
+		ooxmlContentTypesName: docxContentTypes("application/vnd.openxmlformats-officedocument.wordprocessingml.document.main+xml"),
+		"word/document.xml":   "<document/>",
+	})
+	digest := sha256.Sum256(source)
+	metadata := document.AuthorizedUploadMetadata{
+		Filename: "document.docx", MediaFamily: "word", MediaType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+		ByteLength: int64(len(source)), SHA256: hex.EncodeToString(digest[:]),
+		CapabilityRecordChecksum: strings.Repeat("2", 64), ProviderMetadataChecksum: strings.Repeat("3", 64),
+		InputKind: document.RenditionInputOriginalFile,
+	}
+	authorization := document.RenditionAuthorization{
+		ProviderID: descriptor.ID, DescriptorFingerprint: descriptor.Fingerprint, PolicyFingerprint: descriptor.PolicyFingerprint,
+		RenditionRequestFingerprint: strings.Repeat("4", 64), SourceSHA256: metadata.SHA256, SourceBytes: metadata.ByteLength,
+		CapabilityRecordChecksum: metadata.CapabilityRecordChecksum, ProviderMetadataChecksum: metadata.ProviderMetadataChecksum,
+		MediaFamily: metadata.MediaFamily, MediaType: metadata.MediaType, InputKind: metadata.InputKind,
+		MaxProviderMarkdownBytes: 4096, MaxTotalResultBytes: 32768,
+		AuthorizedAt: time.Now().UTC().Add(-time.Minute).Format("2006-01-02T15:04:05.000000000Z"),
+		ExpiresAt:    time.Now().UTC().Add(10 * time.Minute).Format("2006-01-02T15:04:05.000000000Z"),
+	}
+	client := newRenditionTestClient(t, policy, manifest, descriptor,
+		renditionSecrets{"mistral-ocr": "synthetic-key"}, roundTripFunc(func(request *http.Request) (*http.Response, error) {
+			return &http.Response{
+				StatusCode: http.StatusOK, Header: http.Header{"Content-Type": []string{"application/json"}},
+				Body: io.NopCloser(strings.NewReader(response)), Request: request,
+			}, nil
+		}))
+
+	result, err := client.Render(t.Context(), &renditionUpload{Reader: bytes.NewReader(source), metadata: metadata}, authorization)
+	require.NoError(t, err)
+	assert.Equal(t, int64(2), result.Receipt.Usage.Units)
+	assert.Len(t, result.Evidence.Units, 2)
+}
+
 func TestRenditionClientClassifiesHTTPAndModelFailures(t *testing.T) {
 	for _, testCase := range []struct {
 		name   string

document/mistral/spool.go

@@ -329,16 +329,6 @@ func wrapSpoolIOError(operation string, err error) error {
 	return fmt.Errorf("%w: %w", ErrSpoolUnavailable, operationError)
 }
 
-func countLocalUnits(format CandidateFormat, reader io.ReaderAt, size int64) (int, error) {
-	counter := localUnitCounters[format.ID]
-	if counter == nil {
-		return 0, nil
-	}
-	return counter(reader, size)
-}
-
-var localUnitCounters = map[string]func(io.ReaderAt, int64) (int, error){}
-
 // ScavengeSpoolDirectory removes stale package-created regular files. It
 // leaves unrelated regular files and fails closed on unsafe entries.
 func ScavengeSpoolDirectory(directory string, staleBefore time.Time) (int, error) {

document/mistral/units.go

@@ -0,0 +1,17 @@
+package mistral
+
+import "io"
+
+type localUnitCounter func(io.ReaderAt, int64) (int, error)
+
+// localUnitCounters is the Mistral-owned authority registry. Only formats
+// with provider-authentic unit evidence may be added here.
+var localUnitCounters = map[string]localUnitCounter{}
+
+func countLocalUnits(format CandidateFormat, reader io.ReaderAt, size int64) (int, error) {
+	counter := localUnitCounters[format.ID]
+	if counter == nil {
+		return 0, nil
+	}
+	return counter(reader, size)
+}

document/mistral/units_test.go

@@ -0,0 +1,36 @@
+package mistral
+
+import (
+	"bytes"
+	"testing"
+)
+
+func TestLocalUnitCounterRegistryIsMistralOwnedAndBounded(t *testing.T) {
+	for formatID, counter := range localUnitCounters {
+		if counter == nil {
+			t.Fatalf("localUnitCounters[%q] is nil", formatID)
+		}
+		if _, ok := CandidateFormatByID(formatID); !ok {
+			t.Fatalf("localUnitCounters[%q] has no candidate", formatID)
+		}
+	}
+	for formatID, method := range expectedUnitBounds {
+		if method == UnitBoundLocalExact && localUnitCounters[formatID] == nil {
+			t.Fatalf("local exact format %q has no local counter", formatID)
+		}
+	}
+}
+
+func TestCountLocalUnitsLeavesUnprovedFormatsUnbounded(t *testing.T) {
+	docx, ok := CandidateFormatByID("docx")
+	if !ok {
+		t.Fatal("docx candidate is missing")
+	}
+	units, err := countLocalUnits(docx, bytes.NewReader([]byte("synthetic")), 9)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if units != 0 {
+		t.Fatalf("countLocalUnits(docx) = %d, want 0 for unproved format", units)
+	}
+}