Skip to content

Commit c6bae76

Browse files
jgindinjgindin
committed
Treeshake s2c messages to Kotlin SDK
Implement coordinated component and message pruning via withPruning, propagating allowedMessages from A2uiSchemaManager down to Catalog. Add robust automated unit tests and enable full conformance verification. Port of Python SDK commit 0fd7240
1 parent db47eb6 commit c6bae76

8 files changed

Lines changed: 202 additions & 166 deletions

File tree

‎agent_sdks/conformance/suites/catalog.yaml‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,3 +390,74 @@
390390
schema: {type: object, additionalProperties: true}
391391
expect:
392392
schema: {type: object, additionalProperties: true}
393+
394+
- name: test_with_pruning_common_types_transitive
395+
description: Tests that pruning components preserves transitively reachable common types.
396+
catalog:
397+
version: "0.9"
398+
common_types_schema:
399+
$defs:
400+
TypeForCompA:
401+
type: string
402+
$ref: "#/$defs/SubtypeForA"
403+
TypeForCompB: {type: number}
404+
SubtypeForA: {type: boolean}
405+
catalog_schema:
406+
catalogId: basic
407+
components:
408+
CompA: {$ref: "common_types.json#/$defs/TypeForCompA"}
409+
CompB: {$ref: "common_types.json#/$defs/TypeForCompB"}
410+
action: prune
411+
args:
412+
allowed_components: [CompA]
413+
expect:
414+
common_types_schema:
415+
$defs:
416+
TypeForCompA:
417+
type: string
418+
$ref: "#/$defs/SubtypeForA"
419+
SubtypeForA: {type: boolean}
420+
421+
- name: test_render_as_llm_instructions_drops_common_types_no_defs
422+
description: Tests that common types without $defs are omitted in rendering.
423+
catalog:
424+
version: "0.9"
425+
s2c_schema: {s2c: schema}
426+
common_types_schema: {something: else}
427+
catalog_schema:
428+
$schema: "https://json-schema.org/draft/2020-12/schema"
429+
catalog: schema
430+
catalogId: id_basic
431+
action: render
432+
expect_output: |
433+
---BEGIN A2UI JSON SCHEMA---
434+
435+
### Server To Client Schema:
436+
{"s2c":"schema"}
437+
438+
### Catalog Schema:
439+
{"$schema":"https://json-schema.org/draft/2020-12/schema","catalog":"schema","catalogId":"id_basic"}
440+
441+
---END A2UI JSON SCHEMA---
442+
443+
- name: test_render_as_llm_instructions_drops_common_types_empty_defs
444+
description: Tests that common types with empty $defs are omitted in rendering.
445+
catalog:
446+
version: "0.9"
447+
s2c_schema: {s2c: schema}
448+
common_types_schema: {$defs: {}}
449+
catalog_schema:
450+
$schema: "https://json-schema.org/draft/2020-12/schema"
451+
catalog: schema
452+
catalogId: id_basic
453+
action: render
454+
expect_output: |
455+
---BEGIN A2UI JSON SCHEMA---
456+
457+
### Server To Client Schema:
458+
{"s2c":"schema"}
459+
460+
### Catalog Schema:
461+
{"$schema":"https://json-schema.org/draft/2020-12/schema","catalog":"schema","catalogId":"id_basic"}
462+
463+
---END A2UI JSON SCHEMA---

‎agent_sdks/kotlin/src/main/kotlin/com/google/a2ui/core/InferenceStrategy.kt‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@ interface InferenceStrategy {
3030
* @param uiDescription Optional UI context or descriptive instruction.
3131
* @param clientUiCapabilities Capabilities reported by the client for targeted schema pruning.
3232
* @param allowedComponents A specific list of component IDs allowed for rendering.
33+
* @param allowedMessages A specific list of message IDs allowed for rendering.
3334
* @param includeSchema Whether to embed the A2UI JSON schema directly in the instructions.
3435
* @param includeExamples Whether to embed few-shot examples in the instructions.
3536
* @param validateExamples Whether to preemptively validate loaded examples against the schema.
@@ -41,6 +42,7 @@ interface InferenceStrategy {
4142
uiDescription: String = "",
4243
clientUiCapabilities: kotlinx.serialization.json.JsonObject? = null,
4344
allowedComponents: List<String> = emptyList(),
45+
allowedMessages: List<String> = emptyList(),
4446
includeSchema: Boolean = false,
4547
includeExamples: Boolean = false,
4648
validateExamples: Boolean = false,

‎agent_sdks/kotlin/src/main/kotlin/com/google/a2ui/core/parser/StreamingParser.kt‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -422,6 +422,9 @@ abstract class StreamingParser(
422422
}
423423

424424
protected fun processJsonChunk(chunk: String, messages: MutableList<ResponsePart>) {
425+
if (jsonBuffer.length + chunk.length > MAX_JSON_BUFFER_SIZE) {
426+
throw IllegalArgumentException("A2UI JSON buffer exceeded maximum size limit.")
427+
}
425428
for (i in chunk.indices) {
426429
val char = chunk[i]
427430
var charHandled = false
@@ -544,7 +547,9 @@ abstract class StreamingParser(
544547
}
545548
}
546549

547-
if (braceCount > 0 && char in listOf('"', ':', ',', '}', ']')) {
550+
if (
551+
braceCount > 0 && (char == '"' || char == ':' || char == ',' || char == '}' || char == ']')
552+
) {
548553
sniffMetadata()
549554
}
550555
}
@@ -614,6 +619,7 @@ abstract class StreamingParser(
614619
obj != null && obj["id"]?.jsonPrimitive?.content != null && obj.containsKey("component")
615620
) {
616621
handlePartialComponent(obj, messages)
622+
break
617623
}
618624
} catch (e: Exception) {
619625
logger.warning { e.message }
@@ -639,9 +645,9 @@ abstract class StreamingParser(
639645
try {
640646
obj = Json.parseToJsonElement(fixedFragment) as? JsonObject
641647
} catch (_: Exception) {
642-
var trimmed = rawFragment
643-
while ("," in trimmed) {
644-
trimmed = trimmed.substringBeforeLast(",")
648+
var commaIdx = rawFragment.lastIndexOf(',')
649+
while (commaIdx != -1) {
650+
val trimmed = rawFragment.substring(0, commaIdx)
645651
try {
646652
val fixedTrimmed = fixJson(trimmed)
647653
if (fixedTrimmed.isNotEmpty()) {
@@ -650,8 +656,8 @@ abstract class StreamingParser(
650656
}
651657
} catch (ex: Exception) {
652658
logger.warning { ex.message }
653-
continue
654659
}
660+
commaIdx = rawFragment.lastIndexOf(',', commaIdx - 1)
655661
}
656662
}
657663

@@ -1087,7 +1093,7 @@ abstract class StreamingParser(
10871093
private val PREV_KEY_MATCHES_REGEX = Regex("\"key\"\\s*:\\s*\"([^\"]+)\"")
10881094
private val SURFACE_ID_REGEX = Regex("\"surfaceId\"\\s*:\\s*\"([^\"]+)\"")
10891095
private val ROOT_ID_REGEX = Regex("\"root\"\\s*:\\s*\"([^\"]+)\"")
1090-
internal val JSON_NON_PRETTY = Json { prettyPrint = false }
1096+
private const val MAX_JSON_BUFFER_SIZE = 5 * 1024 * 1024
10911097

10921098
/** Factory method returning a version-specific parser instance. */
10931099
fun create(

‎agent_sdks/kotlin/src/main/kotlin/com/google/a2ui/core/schema/A2uiSchemaManager.kt‎

Lines changed: 7 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -176,13 +176,15 @@ constructor(
176176

177177
/**
178178
* Resolves the desired catalog based on the client capabilities, returning it with pruned unused
179-
* components.
179+
* components and messages.
180180
*/
181181
@JvmOverloads
182182
fun getSelectedCatalog(
183183
clientUiCapabilities: JsonObject? = null,
184184
allowedComponents: List<String> = emptyList(),
185-
): A2uiCatalog = selectCatalog(clientUiCapabilities).withPrunedComponents(allowedComponents)
185+
allowedMessages: List<String> = emptyList(),
186+
): A2uiCatalog =
187+
selectCatalog(clientUiCapabilities).withPruning(allowedComponents, allowedMessages)
186188

187189
/** Renders LLM examples for a given catalog, loaded from its configured examples path. */
188190
@JvmOverloads
@@ -197,6 +199,7 @@ constructor(
197199
uiDescription: String,
198200
clientUiCapabilities: JsonObject?,
199201
allowedComponents: List<String>,
202+
allowedMessages: List<String>,
200203
includeSchema: Boolean,
201204
includeExamples: Boolean,
202205
validateExamples: Boolean,
@@ -212,7 +215,8 @@ constructor(
212215
parts.add("## UI Description:\n$uiDescription")
213216
}
214217

215-
val selectedCatalog = getSelectedCatalog(clientUiCapabilities, allowedComponents)
218+
val selectedCatalog =
219+
getSelectedCatalog(clientUiCapabilities, allowedComponents, allowedMessages)
216220

217221
if (includeSchema) {
218222
parts.add(selectedCatalog.renderAsLlmInstructions())

‎agent_sdks/kotlin/src/main/kotlin/com/google/a2ui/core/schema/Catalog.kt‎

Lines changed: 90 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -72,13 +72,28 @@ data class A2uiCatalog(
7272
}
7373

7474
/**
75-
* Returns a new catalog with only allowed components.
75+
* Returns a new catalog with pruned components and messages.
7676
*
7777
* @param allowedComponents List of component names to include.
78-
* @return A copy of the catalog with only allowed components.
78+
* @param allowedMessages List of message names to include in serverToClientSchema.
79+
* @return A copy of the catalog with pruned components and messages.
7980
*/
80-
fun withPrunedComponents(allowedComponents: List<String>): A2uiCatalog {
81-
if (allowedComponents.isEmpty()) return this.withPrunedCommonTypes()
81+
fun withPruning(
82+
allowedComponents: List<String>? = null,
83+
allowedMessages: List<String>? = null,
84+
): A2uiCatalog {
85+
var catalog = this
86+
if (allowedComponents != null) {
87+
catalog = catalog.withPrunedComponentsInternal(allowedComponents)
88+
}
89+
if (allowedMessages != null) {
90+
catalog = catalog.withPrunedMessages(allowedMessages)
91+
}
92+
return catalog.withPrunedCommonTypes()
93+
}
94+
95+
private fun withPrunedComponentsInternal(allowedComponents: List<String>): A2uiCatalog {
96+
if (allowedComponents.isEmpty()) return this
8297

8398
val schemaCopy = catalogSchema.toMutableMap()
8499

@@ -97,68 +112,114 @@ data class A2uiCatalog(
97112
}
98113
}
99114

100-
return copy(catalogSchema = JsonObject(schemaCopy)).withPrunedCommonTypes()
115+
return copy(catalogSchema = JsonObject(schemaCopy))
116+
}
117+
118+
private fun withPrunedMessages(allowedMessages: List<String>): A2uiCatalog {
119+
if (allowedMessages.isEmpty()) return this
120+
121+
val s2cCopy = serverToClientSchema.toMutableMap()
122+
123+
if (version == A2uiVersion.VERSION_0_8) {
124+
(s2cCopy["properties"] as? JsonObject)?.let { props ->
125+
s2cCopy["properties"] =
126+
pruneDefsByReachability(
127+
defs = props,
128+
rootDefNames = allowedMessages,
129+
internalRefPrefix = "#/properties/",
130+
)
131+
}
132+
} else {
133+
(s2cCopy["oneOf"] as? JsonArray)?.let { oneOf ->
134+
val filteredOneOf =
135+
oneOf.filter { item ->
136+
val ref = (item as? JsonObject)?.get("\$ref")?.jsonPrimitive?.content
137+
ref != null && ref.startsWith("#/\$defs/") && ref.split("/").last() in allowedMessages
138+
}
139+
s2cCopy["oneOf"] = JsonArray(filteredOneOf)
140+
}
141+
142+
(s2cCopy["\$defs"] as? JsonObject)?.let { defs ->
143+
s2cCopy["\$defs"] =
144+
pruneDefsByReachability(
145+
defs = defs,
146+
rootDefNames = allowedMessages,
147+
internalRefPrefix = "#/\$defs/",
148+
)
149+
}
150+
}
151+
152+
return copy(serverToClientSchema = JsonObject(s2cCopy))
101153
}
102154

103155
/** Returns a new catalog with unused common types pruned from the schema. */
104156
private fun withPrunedCommonTypes(): A2uiCatalog {
105157
val defs = commonTypesSchema["\$defs"] as? JsonObject ?: return this
106158
if (defs.isEmpty()) return this
107159

108-
fun collectRefs(element: JsonElement, refs: MutableSet<String>) {
109-
when (element) {
160+
val externalRefs = mutableSetOf<String>()
161+
collectRefs(catalogSchema, externalRefs)
162+
collectRefs(serverToClientSchema, externalRefs)
163+
164+
val prefix = "common_types.json#/\$defs/"
165+
val rootDefs =
166+
externalRefs.mapNotNull { if (it.startsWith(prefix)) it.substring(prefix.length) else null }
167+
168+
val newDefs = pruneDefsByReachability(defs, rootDefs)
169+
val newCommonTypes =
170+
JsonObject(commonTypesSchema.toMutableMap().apply { put("\$defs", newDefs) })
171+
172+
return copy(commonTypesSchema = newCommonTypes)
173+
}
174+
175+
private fun collectRefs(rootElement: JsonElement, refs: MutableSet<String>) {
176+
val stack = ArrayDeque<JsonElement>()
177+
stack.addLast(rootElement)
178+
179+
while (stack.isNotEmpty()) {
180+
when (val element = stack.removeLast()) {
110181
is JsonObject -> {
111182
for ((k, v) in element) {
112183
if (k == "\$ref" && v is JsonPrimitive && v.isString) {
113184
refs.add(v.content)
114185
} else {
115-
collectRefs(v, refs)
186+
stack.addLast(v)
116187
}
117188
}
118189
}
119190
is JsonArray -> {
120191
for (item in element) {
121-
collectRefs(item, refs)
192+
stack.addLast(item)
122193
}
123194
}
124195
else -> {}
125196
}
126197
}
198+
}
127199

200+
private fun pruneDefsByReachability(
201+
defs: JsonObject,
202+
rootDefNames: List<String>,
203+
internalRefPrefix: String = "#/\$defs/",
204+
): JsonObject {
128205
val visitedDefs = mutableSetOf<String>()
129-
val queue = ArrayDeque<String>()
130-
131-
val externalRefs = mutableSetOf<String>()
132-
collectRefs(catalogSchema, externalRefs)
133-
collectRefs(serverToClientSchema, externalRefs)
134-
135-
val prefix = "common_types.json#/\$defs/"
136-
for (ref in externalRefs) {
137-
if (ref.startsWith(prefix)) {
138-
queue.add(ref.substring(prefix.length))
139-
}
140-
}
206+
val queue = ArrayDeque(rootDefNames)
141207

142208
while (queue.isNotEmpty()) {
143209
val defName = queue.removeFirst()
144210
if (defs.containsKey(defName) && visitedDefs.add(defName)) {
145211
val defElement = defs[defName]!!
146212
val internalRefs = mutableSetOf<String>()
147213
collectRefs(defElement, internalRefs)
148-
val internalPrefix = "#/\$defs/"
149214
for (ref in internalRefs) {
150-
if (ref.startsWith(internalPrefix)) {
151-
queue.add(ref.substring(internalPrefix.length))
215+
if (ref.startsWith(internalRefPrefix)) {
216+
queue.add(ref.substring(internalRefPrefix.length))
152217
}
153218
}
154219
}
155220
}
156221

157-
val newDefs = JsonObject(defs.filterKeys { it in visitedDefs })
158-
val newCommonTypes =
159-
JsonObject(commonTypesSchema.toMutableMap().apply { put("\$defs", newDefs) })
160-
161-
return copy(commonTypesSchema = newCommonTypes)
222+
return JsonObject(defs.filterKeys { it in visitedDefs })
162223
}
163224

164225
private fun pruneAnyComponentOneOf(

‎agent_sdks/kotlin/src/test/kotlin/com/google/a2ui/conformance/ConformanceTest.kt‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -243,12 +243,6 @@ class ConformanceTest {
243243
val args = case[ConformanceTestHelper.KEY_ARGS] as? Map<*, *> ?: emptyMap<Any, Any>()
244244

245245
// Filter out non-conformant tests for Kotlin
246-
if (
247-
action == "prune" && (args.containsKey("allowed_messages") || name.contains("common_types"))
248-
) {
249-
println("Skipping non-conformant test (prune messages/common_types): $name")
250-
return@mapNotNull null
251-
}
252246
if (
253247
action == "load" &&
254248
(args[KEY_PATH] as? String)?.let {
@@ -273,13 +267,22 @@ class ConformanceTest {
273267

274268
when (action) {
275269
"prune" -> {
276-
val allowedComponents = args[KEY_ALLOWED_COMPONENTS] as? List<String> ?: emptyList()
277-
val pruned = catalog!!.withPrunedComponents(allowedComponents)
270+
val allowedComponents = args[KEY_ALLOWED_COMPONENTS] as? List<String>
271+
val allowedMessages = args["allowed_messages"] as? List<String>
272+
val pruned = catalog!!.withPruning(allowedComponents, allowedMessages)
278273
val expect = case[ConformanceTestHelper.KEY_EXPECT] as Map<*, *>
279274
if (expect.containsKey(KEY_CATALOG_SCHEMA)) {
280275
val expectSchema = jsonMapper.writeValueAsString(expect[KEY_CATALOG_SCHEMA])
281276
assertEquals(Json.parseToJsonElement(expectSchema), pruned.catalogSchema)
282277
}
278+
if (expect.containsKey("s2c_schema")) {
279+
val expectSchema = jsonMapper.writeValueAsString(expect["s2c_schema"])
280+
assertEquals(Json.parseToJsonElement(expectSchema), pruned.serverToClientSchema)
281+
}
282+
if (expect.containsKey("common_types_schema")) {
283+
val expectSchema = jsonMapper.writeValueAsString(expect["common_types_schema"])
284+
assertEquals(Json.parseToJsonElement(expectSchema), pruned.commonTypesSchema)
285+
}
283286
}
284287
"load" -> {
285288
val path = args[KEY_PATH] as? String

0 commit comments

Comments
 (0)