Feature Request: External Trace ID + Execution Metadata + Replay Mode for Vertex AI Agents

[yuer-dsl]

Is your feature request related to a problem?

When experimenting with Vertex AI Agent Builder + Gemini models,
I found it difficult to achieve:

• trace_id propagation across agent / tool steps
• execution metadata (timestamp, tool calls, input preview)
• deterministic replay of a prior agent run

These features are useful for debugging, CI/CD, QA reproducibility,
and enterprise auditability.

Today, users must manually implement such capabilities.

Describe the solution you'd like

1. Allow injecting an external trace_id at agent entry.
2. Propagate metadata automatically across tool calls.
3. Provide a replay mode to deterministically re-run a conversation
   using the same trace_id and metadata.
4. Expose an execution envelope with per-step metadata.

Describe alternatives you've considered

I built a minimal POC (below) implementing:

• external trace_id injection
• metadata envelope
• replay stub
• agent query wrapper

The POC works, but requires custom code.

Additional Context (Full Minimal POC)

flow.dag.yaml

schema_version: 1.0.0
name: minimal-trace-replay-poc
description: Minimal flow to illustrate trace id + execution metadata + replay mode

inputs:
  user_text:
    type: string
    default: "PromptFlow POC input text."
  replay_mode:
    type: boolean
    default: false
  trace_id:
    type: string
    default: null

outputs:
  final_output:
    type: string
    reference: post_process.output

nodes:
  - name: summarize
    type: llm
    provider: openai
    api: chat
    model: gpt-4.1-mini
    inputs:
      prompt: |
        You are a concise assistant.
        Produce a valid JSON response:
        {"summary": "..."}
        Text:
        {{user_text}}

  - name: validate_structure
    type: python
    source:
      type: code
      path: validate.py
    inputs:
      raw: ${summarize.output}

  - name: post_process
    type: python
    source:
      type: code
      path: post_process.py
    inputs:
      text: ${validate_structure.structured}
      trace_id: {{inputs.trace_id}}

  - name: replay_stub
    type: python
    source:
      type: code
      path: replay_stub.py
    inputs:
      trace_id: ${post_process.trace_id}
      replay_mode: {{inputs.replay_mode}}
      metadata: ${post_process.metadata}

validate.py

import json
from typing import Dict

def main(raw: str) -> Dict:
    try:
        data = json.loads(raw)
        if "summary" not in data:
            raise ValueError("Missing 'summary' key")
        return {"valid": True, "structured": data}
    except Exception as e:
        raise ValueError(f"Structure validation failed: {e}\nRaw: {raw}")

post_process.py

import uuid
from datetime import datetime
from typing import Dict

def main(input: Dict) -> Dict:
    summary = input.get("text", {}).get("summary", "")
    trace_id = input.get("trace_id") or str(uuid.uuid4())

    metadata = {
        "trace_id": trace_id,
        "timestamp": datetime.utcnow().isoformat(),
        "node": "post_process",
        "input_preview": summary[:50],
        "replay_intent": "capture_for_replay"
    }

    return {
        "output": f"[POST-PROCESSED] {summary}",
        "trace_id": trace_id,
        "metadata": metadata
    }

replay_stub.py

from typing import Dict

def main(trace_id: str = None, replay_mode: bool = False, metadata: Dict = None):
    if replay_mode and trace_id:
        return {
            "output": f"[REPLAYED_RUN] trace_id={trace_id}",
            "source": "replay_stub",
            "metadata": metadata
        }
    return {"output": "normal_execution"}

These features would greatly enhance reproducibility and observability for
multi-step flows in real-world scenarios.

Thank you for your time and consideration.

I look forward to your feedback and hope we can discuss the possibilities for implementing these features in Vertex AI.
I would be happy to help with further details or contribute if needed.

Best regards,
yuer（https://github.com/yuer-dsl）