Skip to content

Commit b9ee9db

Browse files
josemontespgjosemontespg
committed
Merge remote-tracking branch 'upstream/main' into add-08-integration-tests
2 parents 0f91c98 + f64ccb6 commit b9ee9db

17 files changed

Lines changed: 5027 additions & 155 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/basic_catalog/BasicCatalogProvider.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ import com.google.a2ui.core.schema.A2uiConstants
2323
import com.google.a2ui.core.schema.A2uiVersion
2424
import com.google.a2ui.core.schema.CatalogConfig
2525
import com.google.a2ui.core.schema.SchemaResourceLoader
26+
import com.google.a2ui.core.schema.resolveExamplesPath
2627
import kotlinx.serialization.json.JsonObject
2728
import kotlinx.serialization.json.JsonPrimitive
2829

@@ -91,6 +92,6 @@ object BasicCatalog {
9192
CatalogConfig(
9293
name = BASIC_CATALOG_NAME,
9394
provider = BundledCatalogProvider(version),
94-
examplesPath = examplesPath,
95+
examplesPath = resolveExamplesPath(examplesPath),
9596
)
9697
}

‎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: 29 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -402,13 +402,14 @@ abstract class StreamingParser(
402402

403403
if (isSu && sid != null) {
404404
if (!seenSu.contains(sid)) {
405-
dedupedMsgs.add(0, m)
405+
dedupedMsgs.add(m)
406406
seenSu.add(sid)
407407
}
408408
} else {
409-
dedupedMsgs.add(0, m)
409+
dedupedMsgs.add(m)
410410
}
411411
}
412+
dedupedMsgs.reverse()
412413
messages[i] = part.copy(a2uiJson = dedupedMsgs)
413414
}
414415

@@ -422,6 +423,9 @@ abstract class StreamingParser(
422423
}
423424

424425
protected fun processJsonChunk(chunk: String, messages: MutableList<ResponsePart>) {
426+
if (jsonBuffer.length + chunk.length > MAX_JSON_BUFFER_SIZE) {
427+
throw IllegalArgumentException("A2UI JSON buffer exceeded maximum size limit.")
428+
}
425429
for (i in chunk.indices) {
426430
val char = chunk[i]
427431
var charHandled = false
@@ -477,16 +481,17 @@ abstract class StreamingParser(
477481
if (braceCount >= 0) {
478482
val objBuffer = jsonBuffer.substring(startIdx)
479483
if (objBuffer.startsWith("{") && objBuffer.endsWith("}")) {
484+
val isTopLevel =
485+
braceStack.isEmpty() ||
486+
(inTopLevelList && braceStack.size == 1 && braceStack[0].first == "[")
487+
480488
try {
481489
val obj = Json.parseToJsonElement(objBuffer) as? JsonObject
482490
if (obj != null) {
483491
foundValidJsonInBlock = true
484492

485493
val isProtocol = inTopLevelList && isProtocolMsg(obj)
486494
val isComp = obj.containsKey("id") && obj.containsKey("component")
487-
val isTopLevel =
488-
braceStack.isEmpty() ||
489-
(inTopLevelList && braceStack.size == 1 && braceStack[0].first == "[")
490495

491496
if (isComp) {
492497
handlePartialComponent(obj, messages)
@@ -513,11 +518,8 @@ abstract class StreamingParser(
513518
}
514519
} catch (e: Exception) {
515520
if (
516-
(e is IllegalArgumentException &&
517-
e !is kotlinx.serialization.SerializationException) ||
518-
e.message?.contains("Circular reference") == true ||
519-
e.message?.contains("Self-reference") == true ||
520-
e.message?.contains("Validation failed") == true
521+
e is IllegalArgumentException &&
522+
e !is kotlinx.serialization.SerializationException
521523
) {
522524
throw e
523525
}
@@ -544,7 +546,9 @@ abstract class StreamingParser(
544546
}
545547
}
546548

547-
if (braceCount > 0 && char in listOf('"', ':', ',', '}', ']')) {
549+
if (
550+
braceCount > 0 && (char == '"' || char == ':' || char == ',' || char == '}' || char == ']')
551+
) {
548552
sniffMetadata()
549553
}
550554
}
@@ -583,7 +587,9 @@ abstract class StreamingParser(
583587
"root" -> ROOT_ID_REGEX.find(jsonBuffer, idx)
584588
else -> {
585589
val fragment = jsonBuffer.substring(idx)
586-
Regex("\"$key\"\\s*:\\s*\"([^\"]+)\"").find(fragment)
590+
val regex =
591+
LATEST_VALUE_REGEX_CACHE.getOrPut(key) { Regex("\"$key\"\\s*:\\s*\"([^\"]+)\"") }
592+
regex.find(fragment)
587593
}
588594
}
589595
if (match != null) {
@@ -614,6 +620,7 @@ abstract class StreamingParser(
614620
obj != null && obj["id"]?.jsonPrimitive?.content != null && obj.containsKey("component")
615621
) {
616622
handlePartialComponent(obj, messages)
623+
break
617624
}
618625
} catch (e: Exception) {
619626
logger.warning { e.message }
@@ -624,6 +631,7 @@ abstract class StreamingParser(
624631

625632
protected fun sniffPartialDataModel(messages: MutableList<ResponsePart>) {
626633
val msgType = dataModelMsgType
634+
627635
if (jsonBuffer.indexOf("\"$msgType\"") == -1) return
628636

629637
for (i in braceStack.indices.reversed()) {
@@ -639,9 +647,9 @@ abstract class StreamingParser(
639647
try {
640648
obj = Json.parseToJsonElement(fixedFragment) as? JsonObject
641649
} catch (_: Exception) {
642-
var trimmed = rawFragment
643-
while ("," in trimmed) {
644-
trimmed = trimmed.substringBeforeLast(",")
650+
var commaIdx = rawFragment.lastIndexOf(',')
651+
while (commaIdx != -1) {
652+
val trimmed = rawFragment.substring(0, commaIdx)
645653
try {
646654
val fixedTrimmed = fixJson(trimmed)
647655
if (fixedTrimmed.isNotEmpty()) {
@@ -650,8 +658,8 @@ abstract class StreamingParser(
650658
}
651659
} catch (ex: Exception) {
652660
logger.warning { ex.message }
653-
continue
654661
}
662+
commaIdx = rawFragment.lastIndexOf(',', commaIdx - 1)
655663
}
656664
}
657665

@@ -990,6 +998,9 @@ abstract class StreamingParser(
990998
if (pathElem != null) {
991999
val currentPath = pathElem.jsonPrimitive.content
9921000
if (!currentPath.startsWith("/")) {
1001+
if (!map.containsKey("componentId")) {
1002+
map.clear()
1003+
}
9931004
map["path"] = JsonPrimitive("/$currentPath")
9941005
}
9951006
}
@@ -1087,7 +1098,8 @@ abstract class StreamingParser(
10871098
private val PREV_KEY_MATCHES_REGEX = Regex("\"key\"\\s*:\\s*\"([^\"]+)\"")
10881099
private val SURFACE_ID_REGEX = Regex("\"surfaceId\"\\s*:\\s*\"([^\"]+)\"")
10891100
private val ROOT_ID_REGEX = Regex("\"root\"\\s*:\\s*\"([^\"]+)\"")
1090-
internal val JSON_NON_PRETTY = Json { prettyPrint = false }
1101+
private val LATEST_VALUE_REGEX_CACHE = mutableMapOf<String, Regex>()
1102+
private const val MAX_JSON_BUFFER_SIZE = 5 * 1024 * 1024
10911103

10921104
/** Factory method returning a version-specific parser instance. */
10931105
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())

0 commit comments

Comments
 (0)