[api][runtime][python] Propagate metric groups to cross-language resources

[goutamadwant]

Linked issue: #857

Purpose of change

This change propagates the action metric group across Java/Python resource boundaries so cross-language resources record metrics under the caller's action scope.

The forwarding is handled by the Python-backed wrapper classes rather than by adding Python-specific logic to the base Resource class. On the Python-to-Java path, Java resource wrappers now unwrap only FlinkMetricGroup instances or None, and fail clearly for unsupported custom metric groups instead of passing an opaque Python object through pemja.

Metric ownership remains single-recorder: this PR forwards the bound metric group but does not add provider-side token recording. It also keeps the return-path/effective-scope question out of scope. If a provider/resource wraps the metric group with provider-specific dimensions before action-level token recording, that behavior should be tracked separately.

Tests

• uv run --no-sync pytest flink_agents/runtime/tests/test_cross_language_metric_group.py
• mvn --batch-mode --no-transfer-progress -pl api,plan,runtime -am -DskipTests -Dspotless.skip=true test-compile
• mvn --batch-mode --no-transfer-progress -pl api -Dtest=PythonChatModelSetupTest -Dspotless.skip=true test
• mvn --batch-mode --no-transfer-progress -pl runtime -am -Dtest=PythonResourceAdapterImplTest -Dsurefire.failIfNoSpecifiedTests=false -Dspotless.skip=true test
• ./tools/lint.sh -c

API

Touches the Java/Python bridge interfaces by adding default metric-forwarding hooks on PythonResourceAdapter and PythonResourceWrapper. There are no user-facing configuration changes.

Documentation

• doc-needed
• doc-not-needed
• doc-included

[weiqingy]

Thanks for taking this on — a few questions inline.

[weiqingy]

This now forwards the action metric group through to the Python provider, which is exactly what issue #857 asked for — and #857 also called out "define metric ownership clearly to avoid double-counting when action-level and provider-level metrics both observe the same response." From what I can see, ownership already looks singular: token metrics are recorded in exactly one place per language — the action handler (ChatModelAction.recordChatTokenMetrics → chatModel.recordTokenMetrics on the Java side, chat_model_action → _record_token_metrics on the Python side), each reading the resource's bound group. The provider integrations under integrations/ record nothing, so there's no second recorder to collide with — this PR widens that single existing recording across the language boundary rather than introducing a second one. So the #857 concern reads as already closed by construction. Would it be worth a sentence in the PR description making that explicit, so the "no double-counting" reasoning is traceable back to the issue?

[joeyutong]

@weiqingy One ownership detail I think is still worth calling out: this is not only about avoiding a second recorder after forwarding the metric group. It is also about which effective metric group the action-level recorder should use after the resource call returns.

For example, in the Java action -> Python chat model path, the Python resource may wrap the forwarded action metric group with a provider-specific dimension such as modelProvider. Metrics recorded inside the Python resource would use that wrapped group, but token metrics are still recorded later by the Java action/model-setup path, which only sees the Java wrapper's metric group. That means action-recorded token metrics may still miss the provider-specific dimension.

So I think the ownership question also needs to define whether a resource-derived metric scope should flow back to the action-level recorder. This PR fixes forwarding into the resource, but not that return path.

[weiqingy]

Thanks, that makes sense. My earlier comment was focused on the double-recording angle, but I agree there is a separate ownership question around the effective metric scope.

My takeaway is that #860 currently fixes the forward propagation path into the cross-language resource, but it does not settle whether a resource/provider-derived metric scope should flow back to the action-level token recorder. For example, if the Python resource wraps the forwarded group with provider-specific dimensions, Java-side token recording after the call may still use only the Java wrapper/action scope.

If we keep #860 scoped to forward propagation, it may be good to call this out in the PR description or track it as a follow-up, so it is clear that the return-path/effective-scope behavior is still unresolved. Does that match what you had in mind?

[joeyutong]

Thanks, that matches my understanding.

I agree that #860 can stay focused on the forward propagation path, and the return-path / effective-scope behavior can be tracked separately.

@goutamadwant Since this PR already calls out that limitation, would you mind filing a follow-up issue for it? That would make the remaining ownership question easier to continue without blocking this PR.

[weiqingy]

The abstract base Resource (package ...api.resource) now imports and instanceof-checks PythonResourceWrapper from the more specialized ...api.resource.python sub-package. That inverts the usual dependency direction — every Java-only resource (chat models, tools, prompts, vector stores) now carries a compile-time edge to the Python-bridge interface in its base type, and a hypothetical third forwarding flavor would mean editing this base method again rather than overriding it.

One option that keeps the base oblivious to the bridge: have each Python* wrapper override setMetricGroup to call super.setMetricGroup(...) and then setPythonResourceMetricGroup(...), pushing the Python concern down into the wrappers where getPythonResourceAdapter() already lives. The PR already adds a new getPythonResourceAdapter() override to all eight wrappers, so a per-wrapper setMetricGroup override would touch the same set of files — just in the Python layer instead of the base class. Was the centralized instanceof chosen deliberately, or mainly to avoid touching the wrappers?

[weiqingy]

This new file covers the Python→Java direction well, but only that direction — set_java_resource_metric_group plus the four Java*Impl wrappers. The Java→Python direction's Python half, python_java_utils.set_metric_group (python_java_utils.py:382), has no direct test: it's the function that wraps the raw Java group in FlinkMetricGroup and owns the only None branch in the feature, and the Java half is mock-verify only (testSetMetricGroup asserts the invoke string, it doesn't execute Python). So that seam isn't exercised end-to-end today.

Would a small unit test in this file's style be worth adding — a fake resource that captures its set_metric_group arg, then set_metric_group(fake, sentinel_j_group) asserting the captured arg is a FlinkMetricGroup whose _j_metric_group is sentinel_j_group, plus a None case asserting None is forwarded? That would pin the wrap-and-None behavior that's currently uncovered.

[weiqingy]

In the real flow this is correct — metric_group is always a FlinkMetricGroup (it comes from action_metric_group), so the unwrap hits _j_metric_group and forwards the genuine Java group. Low confidence that the edge below matters in practice, so genuinely a question rather than a flag.

RunnerContext.get_resource(name, type, metric_group=...) is a documented public API accepting any MetricGroup — an abstract base a user could subclass without a _j_metric_group. If such a custom group reaches a Java resource wrapper, the silent , metric_group fallback would forward a raw Python object into Java setMetricGroup(FlinkAgentsMetricGroup), failing opaquely at the pemja boundary rather than with a clear error. Worth making the fallback explicit (unwrap only a real FlinkMetricGroup, else fail loudly), or is that out of scope today given no in-tree caller hits it? (The None case is benign — get_resource substitutes action_metric_group first, and a direct None just clears the field — so no concern there.)

[goutamadwant]

@weiqingy @joeyutong Pushed an update addressing the review comments.

The base Resource class no longer depends on the Python bridge. Metric forwarding now lives in the Python-backed wrappers, which call super.setMetricGroup(...) and then forward to the wrapped Python resource.

I also added direct coverage for python_java_utils.set_metric_group(...), including the raw Java group wrapping path and the None path. On the Python-to-Java side, set_java_resource_metric_group(...) now accepts only FlinkMetricGroup or None and raises a clear TypeError for unsupported custom metric groups.

I also updated the PR description to make the scope explicit: this PR handles forward metric-group propagation and keeps token recording single-recorder/no double-counting. The return-path/effective-scope behavior is still separate follow-up work. Let me know if you have any suggestions. thanks!