Add Azure AI Emulation (CognitiveServices + Machine Learning, ARM + Data Plane)#232
Conversation
Add the azureai driver interfaces and in-memory provider Mock spanning both ARM providers that make up Azure AI — Microsoft.CognitiveServices (AI Foundry / AI Studio / the AI Services resource / Azure OpenAI) and Microsoft.MachineLearningServices (Azure ML) — plus the data planes: Azure OpenAI inference (chat/completions, embeddings, completions), the AI Foundry Agents/Assistants API, and AML online-endpoint scoring. CognitiveServices covers accounts (CRUD, keys, models, skus, usages), deployments, projects, raiPolicies, commitmentPlans, and private-endpoint connections. MachineLearningServices covers workspaces, computes (with a start/stop/restart state machine), online/batch endpoints + deployments, jobs, versioned assets (models/data/environments/components/featuresets), datastores, connections, schedules, and registries. Slices/maps are deep-copied on store and return and mutators use copy-then-Set. Auto-metrics push to Azure Monitor via SetMonitoring; wired into the Azure provider.
Serve Microsoft.CognitiveServices and Microsoft.MachineLearningServices over
the ARM JSON wire protocol, plus the Azure OpenAI inference / Assistants and
AML scoring data planes. Real armcognitiveservices / armmachinelearning clients
(and any REST caller) configured with a custom endpoint hit these handlers the
same way they hit management.azure.com.
CognitiveServices uses the shared azurearm.ParsePath; Azure ML nests deeper
(onlineEndpoints/{e}/deployments/{d}, models/{n}/versions/{v}, computes/{c}/start)
so its handler parses the trailing segments itself. The data plane is
host/path-routed on /openai/ and /score, with the account scope derived from the
request Host subdomain. Roundtrip tests start a real httptest server and cover
every family over HTTP. Wired into the Azure server Drivers bundle.
Add the portable azureai.AzureAI Layer-1 wrapper implementing the full driver.AzureAI surface through a single do() pipeline (recorder, metrics, rate limiting, error injection, latency), mirroring sagemaker/vertexai. Add a chaos wrapper (chaos.WrapAzureAI) that injects failures on account/workspace and deployment/job creation and the inference/scoring runtimes, plus default Azure AI cost rates. Brings Azure AI to full parity with how the other AI services participate in the cross-cutting layers.
…i-sdk-compat # Conflicts: # cost/cost.go # docs/services.md
Add the azuresearch driver interfaces and in-memory provider Mock for Azure
AI Search (Microsoft.Search) — the ARM control plane (searchServices lifecycle,
admin/query keys, shared private links, private-endpoint connections) and the
{service}.search.windows.net data plane (indexes, documents with
upload/merge/delete + search/suggest/autocomplete/count/get, indexers with
run/reset/status, data sources, skillsets, synonym maps, aliases, service
statistics). Slices/maps are deep-copied and mutators copy-then-Set; document
search is a deterministic substring match bounded by maxSearchTop. Auto-metrics
push to Azure Monitor; wired into the Azure provider.
Serve Microsoft.Search/searchServices over ARM and the search data plane over
its host/path-routed surface (/indexes, /indexes/{i}/docs/*, /indexers,
/datasources, /skillsets, /synonymmaps, /aliases, /servicestats). The data-plane
service scope is derived from the {service}.search.windows.net Host subdomain.
Roundtrip tests start a real httptest server and cover service lifecycle + keys
+ private links and the full index/document/indexer/datasource/skillset/synonym/
alias/stats surface. Wired into the Azure server Drivers bundle.
Add the portable azuresearch.AzureSearch Layer-1 wrapper implementing the full driver.AzureSearch surface through a do() pipeline (recorder, metrics, rate limiting, error injection, latency), a chaos wrapper (chaos.WrapAzureSearch) that injects failures on service/index creation, document indexing and the query runtime, and default Azure AI Search cost rates.
NitinKumar004
left a comment
NitinKumar004
left a comment
There was a problem hiding this comment.
Review: Azure AI + Azure AI Search SDK-compat
Great scope — two whole services (Cognitive Services + Azure OpenAI + Machine Learning, and AI Search) with control-plane + data-plane. It builds cleanly and the full go test ./... tree is green. I ran a multi-angle review (both server handlers, both provider mocks, and the wiring/chaos/cost surface) and verified findings against source. Dispatch/registration is correct — the ARM handlers are provider-scoped and the data-plane paths are disjoint from Databricks/Cosmos/Functions (one caveat below).
The headline: these aren't real-SDK roundtrip tests
Despite the *_roundtrip_test.go / sdk_roundtrip_test.go names, none of these tests import the real Azure SDKs (azsearch, armcognitiveservices, armmachinelearning, azopenai). They issue raw http.NewRequest calls with hand-built URLs and decode into map[string]any. Every other SDK-compat slice in this repo (Secrets, DNS, RDS, …) drives the real client, which is what actually proves wire compatibility. Because these don't, the contract bugs below ship green. Strong recommendation: add at least one real-SDK roundtrip per surface — it would have caught most of the inline findings.
Inline findings
Marked inline. The most impactful: mergeOrUpload drops fields, key rotation is a no-op in both services, $filter is ignored, AML /score collapses all endpoints to default, and the data-plane error code is numeric where the SDK expects a string.
Additional items (not inline)
- Real-SDK OData forms untested:
GET /indexes/{i}/docs('{key}')and/indexes('{name}')(paren form the SDK emits) won't route —dataplane.gomatchesparts[2] == "docs"/ exact first-segment roots, so the paren form falls through. Untested because the tests use the slash form. - Blob shadowing (narrow): the Search data-plane
Matchesclaims bare first segments (indexes,aliases,datasources, …); when both Search and Blob drivers are enabled, a blob container namedaliases/indexesis captured by the Search handler. embeddingsdrops non-string inputs: token-array inputs ([[1,2,3]], valid per the API) silently yield zero rows.- Loose action routing:
/openai/deployments/{d}/chat/anythingserves a chat completion (no validation of the trailingcompletionssegment); unknown deployments aren't 404'd. - Chaos wrappers are selective:
WrapAzureSearch/WrapAzureAIonly intercept a few ops, soServiceOutage("azuresearch")leaves reads/indexer/runtime ops live — confirm that's intentional (differs fromwrappers_dns.go, which wraps every method). - Reuse: the five near-identical
serve{Indexers,DataSources,Skillsets,SynonymMaps,Aliases}handlers could share a genericserveCollection[T]helper.
Overall: solid foundation, but I'd address the correctness findings and add real-SDK tests before merge. None of these are blockers to the architecture — they're fidelity/semantics fixes.
| } | ||
|
|
||
| m.documents.Set(storeKey, merged) | ||
| default: // upload | mergeOrUpload |
There was a problem hiding this comment.
mergeOrUpload silently overwrites instead of merging. This branch handles upload and mergeOrUpload as a full replace via documents.Set(...). For mergeOrUpload/merge on an existing doc, Azure preserves fields not present in the payload; here {id:1,name:"x",price:9} + mergeOrUpload {id:1,price:12} becomes {id:1,price:12} — name is dropped. Merge into the existing stored doc for the merge actions.
| matched := make([]driver.SearchResult, 0, len(docs)) | ||
|
|
||
| for _, d := range docs { | ||
| if term == "" || term == "*" || docContains(d, term) { |
There was a problem hiding this comment.
SearchDocuments ignores Filter/OrderBy/Top/Skip. The match loop only tests free-text term; req.Filter is never consulted (grep finds no reference). A query with filter: "price gt 100" returns every doc and a count over the full set. The GET search path (dataplane_docs.go) also drops $top/$skip/$orderby/$select.
| Secondary: hashHex(resourceGroup, name, "admin-secondary"), | ||
| } | ||
|
|
||
| salt := "regen-" + m.now() |
There was a problem hiding this comment.
Admin-key regeneration is a no-op. The salted key is returned once but never persisted; ListAdminKeys always recomputes the unsalted deterministic value, so RegenerateAdminKey → ListAdminKeys returns the original key. Key rotation can't be observed. (Same pattern for query keys, and for Cognitive Services in providers/azure/azureai/azureai.go:301.)
| results := make([]driver.IndexResult, 0, len(actions)) | ||
|
|
||
| for _, act := range actions { | ||
| docKey := fmt.Sprintf("%v", act.Document[kf]) |
There was a problem hiding this comment.
Keyless documents collide + delete-missing reports success. docKey := fmt.Sprintf("%v", act.Document[kf]) yields "<nil>" when the key field is absent, so every keyless doc maps to one store key (second overwrites first, CountDocuments→1). And a delete for a never-indexed key still returns Status:true/200 (memstore miss ignored), unlike the merge branch which correctly 404s.
|
|
||
| salt := "regen-" + m.now() | ||
| if strings.EqualFold(keyName, "Key2") { | ||
| keys.Key2 = deterministicKey(resourceGroup, name, salt) |
There was a problem hiding this comment.
RegenerateAccountKey is a no-op. Keys are recomputed deterministically from (rg,name) with no key store, so the regenerated value is discarded on return; a subsequent ListAccountKeys returns the original key. Rotation round-trip is broken (mirrors the Search admin-key issue).
|
|
||
| for k, a := range m.assets.All() { | ||
| if strings.HasPrefix(k, prefix) { | ||
| if cur, ok := latest[a.Name]; !ok || a.Version > cur.Version { |
There was a problem hiding this comment.
"Latest" asset version chosen by lexical string compare. a.Version > cur.Version compares raw strings, so "10" > "9" is false — registering versions 1..10 reports 9 as latest. Compare numerically when the versions are integers.
| return | ||
| } | ||
|
|
||
| out, cerr := h.dp.ScoreOnlineEndpoint(r.Context(), accountFromHost(r.Host), body) |
There was a problem hiding this comment.
AML /score collapses every online endpoint to "default". The endpoint name comes from accountFromHost, which only recognizes .openai./.services.ai./.cognitiveservices. infixes — not the real AML host {endpoint}.{region}.inference.ml.azure.com. So ep1, ep2, … all map to "default" and endpoint identity/metrics are lost.
| func dpErr(w http.ResponseWriter, status int, msg string) { | ||
| w.Header().Set("Content-Type", "application/json") | ||
| w.WriteHeader(status) | ||
| _ = json.NewEncoder(w).Encode(map[string]any{"error": map[string]any{"message": msg, "code": status}}) |
There was a problem hiding this comment.
Data-plane error envelope emits a numeric code. "code": status writes the HTTP status int, but Azure OpenAI / azcore.ResponseError treat error.code as a string identifier (e.g. "DeploymentNotFound"). A client matching on the string code mis-parses the error. (Same convention in azuresearch/dataplane.go.)
|
|
||
| switch rp.SubResource { | ||
| case subAdminKeys: | ||
| keys, err := h.svc.ListAdminKeys(r.Context(), rg, name) |
There was a problem hiding this comment.
Key actions respond to any HTTP method. serveServiceAction never checks r.Method, so GET .../searchServices/{name}/listAdminKeys returns primaryKey/secondaryKey (real Azure requires POST), and a stray GET .../deleteQueryKey/{k} deletes a key. The control serveService path does switch on method — mirror that here. (Azure AI serveAccountAction has the same gap.)
| }) | ||
| } | ||
|
|
||
| dpJSON(w, map[string]any{"value": out}) |
There was a problem hiding this comment.
Partial-failure document batches always return HTTP 200. The provider correctly emits per-doc Status:false/404, but Azure returns 207 Multi-Status on partial success and the azsearch SDK keys its IndexBatchException off the 207. Because this always writes 200, callers never see partial-batch failures.
Resolve PR stackshy#232 review feedback across the CognitiveServices, MachineLearningServices, and Azure AI Search provider and server layers: - Search: mergeOrUpload now merges onto the existing document instead of replacing it; reject documents missing the key field; honor $filter (single OData comparison), $orderby, $top, $skip and $select; return 207 Multi-Status on partial-failure index batches. - Persist regenerated admin/query and Cognitive Services account keys so rotation is observable via a subsequent list. - Sort ListMessages/ListAssistants by creation sequence for deterministic ordering; pick the latest asset version numerically, not lexically. - Derive the AML online-endpoint name from the inference host so /score no longer collapses every endpoint to "default". - Emit a string error.code (not the numeric status) in data-plane error envelopes, and enforce the correct HTTP method on key/account actions. Add regression tests to the existing roundtrip test files.
NitinKumar004
left a comment
NitinKumar004
left a comment
There was a problem hiding this comment.
Re-review after 891af743a — all 11 inline findings resolved ✅
Thanks for the thorough pass. I re-verified each finding against the current head and the whole tree still builds green with both packages' tests passing. Every concrete correctness bug I flagged is fixed:
| Finding | Fix |
|---|---|
mergeOrUpload dropped fields |
now merges via mergeInto, preserving unspecified fields |
$filter/$orderby/$top/$skip/$select ignored |
real OData filter parser + orderResults + paginate; plumbed through both the GET and POST search paths |
| admin/query & account key regen were no-ops | persisted to a key store, salted with a monotonic seq — rotation is now observable via a subsequent list |
| keyless docs collided | keyless/empty key now returns 400/Status:false instead of colliding on "<nil>" |
ListMessages random order |
sorted by the monotonic sequence encoded in the id |
| "latest" asset version by lexical compare | versionNewer compares numerically |
AML /score collapsed to default |
endpointFromHost parses *.inference.ml.azure.com |
numeric data-plane error code |
codeForStatus() returns string identifiers |
| key actions accepted any method | method-map + 405 on both serveServiceAction and serveAccountAction |
| partial-batch always 200 | emits 207 Multi-Status when any doc fails |
One thing I got wrong in the original review: delete is idempotent in Azure AI Search — returning 200 for a missing key is correct, so keeping that behavior (with the added comment) is the right call. 👍
Still open
The main recommendation — real-SDK roundtrip tests — is still unaddressed. The *_roundtrip_test.go files continue to use hand-built http.NewRequest + map[string]any decoding; none import armcognitiveservices / armmachinelearning / azopenai / azsearch. The new cases are good regression coverage for the fixes above, but without at least one real-client roundtrip per surface we still can't prove wire compatibility — and it's what would settle the remaining items below. Every other SDK-compat slice in the repo drives the real client, so this is also a consistency gap.
Smaller items still open (all non-blocking):
- Embeddings drop non-string input:
parseEmbeddingInputonly handles string /[]string; a token-array input ([[1,2,3]], valid per the API) still yields zero rows. - OData paren-form routing:
docKey()now strips a standalone('key')segment, but the fused forms the SDK emits —/indexes('{name}')and/docs('{key}')— still don't route (Matcheskeys on the bare first segment;serveDocsrequiresparts[2] == "docs"exactly). This is exactly what a real-SDK test would surface. - Blob shadowing (narrow): the Search data-plane
Matchesstill claims bare first segments (indexes,aliases,datasources, …), so a blob container with one of those names is captured by the Search handler when both drivers are enabled. - Chaos wrappers selective:
WrapAzureSearchintercepts 4 ops andWrapAzureAI8, soServiceOutage(...)leaves reads/indexer/runtime ops live — fine if intentional, just flagging it differs from e.g.wrappers_dns.gowhich wraps every method.
Net: the correctness fixes are solid and verified. I'd still like a real-SDK roundtrip per surface before merge — it closes the loop on the paren-form/embeddings/host-parsing fidelity that hand-rolled HTTP tests can't.
Address the re-review follow-ups on PR stackshy#232: - Embeddings: parseEmbeddingInput now accepts token-array inputs ([]int and [][]int) in addition to string / []string, rendering each to a synthetic string so every input yields one embedding row. - Search data plane: normalize the fused OData key form the SDKs emit, so /indexes('name') and /docs('key') route identically to the slash forms. Add real-client roundtrip tests driving the official Azure SDKs against the in-memory server, proving wire compatibility per surface: - armcognitiveservices: account lifecycle + key rotation - armmachinelearning: workspace + versioned assets (numeric latest version) - azopenai: chat completions + embeddings - armsearch: service lifecycle + admin/query key rotation Azure ships no Go data-plane search SDK, so the search data plane's wire fidelity (paren-form routing) is covered by explicit fused-URL tests.
|
Pushed Concrete bugs
Real-SDK roundtrip tests (the main recommendation)Added real-client roundtrips driving the official Azure SDKs against the in-memory server, one per available surface:
One honest caveat: Azure ships no Go data-plane search SDK — Intentionally left as-is (non-blocking, per your flags)
|
…i-sdk-compat # Conflicts: # go.mod # server/azure/azure.go
NitinKumar004
left a comment
NitinKumar004
left a comment
There was a problem hiding this comment.
Approving ✅
Re-reviewed at the current head (merge-current with development). All prior findings are resolved and verified:
- All 11 original correctness findings — fixed and confirmed against source.
- Real-SDK roundtrip tests — the headline ask is done:
azopenai,armcognitiveservices,armmachinelearning/v4, andarmsearchnow drive the in-memory server end-to-end, usingBeginCreate/PollUntilDone(proving the ARM LRO/provisioning-state responses satisfy the real pollers) and asserting key rotation withNotEqual+ a persistence re-read. - Embeddings now accept token-array inputs (
[]int/[][]int); OData paren-form routing (indexes('x'),docs('k')) is normalized and covered by a dedicated test. - Merge integrated cleanly:
go build ./...and the fullgo test ./...tree are green (185 packages), and the Azure AI/Search suites pass on the merged head.
Great work turning this around.
Non-blocking nit (fix on the way in, or as a follow-up)
golangci-lint (repo config enables prealloc) flags 3 spots in providers/azure/azuresearch/dataplane_more.go (lines 57, 133, 177) — out := make([]T, 0) should be make([]T, 0, len(childKeysUnder(...))). Trivial, but our merge gate is 0 lint issues.
Leaving as-is (acknowledged, minor)
- Search data-plane
Matchesclaims bare roots (indexes,aliases, …) — narrow blob-shadow overlap only when both drivers are enabled. - Chaos wrappers are selective (Search 4 ops, AI 8) — fine if intentional.
2 participants
Summary
Adds full Azure AI emulation across all three layers, spanning both ARM providers that make up "Azure AI" —
Microsoft.CognitiveServices(Azure AI Foundry / AI Studio, the AI Services resource, Azure OpenAI) andMicrosoft.MachineLearningServices(Azure Machine Learning) — plus the data planes. Third AI platform after AWS SageMaker (#222) and GCP Vertex AI (#231).azureai/driver,providers/azure/azureai) — 92 operations: 31 CognitiveServices + 46 MachineLearningServices + 15 data plane.server/azure/azureai) — realarmcognitiveservices/armmachinelearningclients (or any REST caller) on a custom endpoint hit these the same way they hitmanagement.azure.com. ARM PUT returns the resource inline with a terminalprovisioningState.azureai/azureai.go) — recorder / metrics / rate-limit / error-injection / latencydo()pipeline, mirroringsagemaker/vertexai.chaos.WrapAzureAI, default cost rates, auto-metrics to Azure Monitor, lifecycle state machines (compute start/stop, terminal provisioning states).Coverage
CognitiveServices: accounts (CRUD, list by RG/sub, listKeys, regenerateKey, listModels, listSkus, listUsages), deployments, projects (AI Foundry), raiPolicies, commitmentPlans, privateEndpointConnections.
Data plane: Azure OpenAI
chat/completions·completions·embeddings; AI Foundry Agents/Assistants (assistants, threads, messages, runs); AML/score.MachineLearningServices: workspaces (Default/Hub/Project/FeatureStore), computes (+start/stop/restart), online/batch endpoints + deployments, jobs (+cancel), versioned assets (models/data/environments/components/featuresets), datastores, connections, schedules, registries.
Notes (honesty flags)
azurearm.ParsePathonly reaches{name}/{sub}/{subName}; AML nests deeper (onlineEndpoints/{e}/deployments/{d},models/{n}/versions/{v},computes/{c}/start), so the AML handler parses trailing segments itself.Test plan
go build ./...go test ./...— all greengolangci-lint runclean on every new package (driver, provider, servers, Layer-1, chaos, cost)httptestserver: account lifecycle, keys/catalogs, deployments/projects/raiPolicies/commitmentPlans/PEC; AML workspaces, compute lifecycle, online endpoints + deployments, jobs + assets, datastores/connections/schedules/registries/scoreUpdate: Azure AI Search added (
Microsoft.Search)This PR now also adds Azure AI Search (the RAG/retrieval backbone) as its own
azuresearchpackage — a distinct ARM provider — with full parity (53 operations: 19 control + 34 data plane).Microsoft.Search/searchServices): service CRUD, list by RG/sub, update; admin keys (list/regenerate), query keys (list/create/delete); shared private links + private-endpoint connections.{service}.search.windows.net): indexes; documents (upload/merge/delete batch, search +count, suggest, autocomplete, count, get-by-key); indexers (+run/reset/status); data sources; skillsets; synonym maps; aliases; service statistics.azuresearch/azuresearch.go),chaos.WrapAzureSearch, cost rates, Azure Monitor auto-metrics.httptestserver cover the control plane (lifecycle, keys, private links) and the entire data plane.Notes: document search is a deterministic substring match (bounded by
maxSearchTop) suitable for emulation, not a real scoring engine. Data-plane service scope is derived from the request Host subdomain.