Repository: pydantic/pydantic-ai
File: tests/test_temporal.py (Python)

Code before:
        workflows=[ImageAgentWorkflow],
        plugins=[AgentPlugin(image_temporal_agent)],
    ):
        with workflow_raises(
            UserError,
            snapshot('Image output is not supported with Temporal because of the 2MB payload size limit.'),
        ):
            await client.execute_workflow(
                ImageAgentWorkflow.run,
                args=['Generate an image of an axolotl.'],
                id=ImageAgentWorkflow.__name__,
                task_queue=TASK_QUEUE,
            )


# ============================================================================
# DocumentUrl Serialization Test - Verifies that DocumentUrl with custom
# media_type is properly serialized through Temporal activities
# ============================================================================


def document_url_response(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
    """Return a DocumentUrl with a custom media_type that cannot be inferred from the URL."""
    assert info.output_tools is not None
    document_url_json = '{"url": "https://example.com/doc/12345", "media_type": "application/pdf"}'
    return ModelResponse(parts=[ToolCallPart(info.output_tools[0].name, document_url_json)])


document_url_model = FunctionModel(document_url_response)

document_url_agent = Agent(
    document_url_model,
    name='document_url_agent',
    output_type=DocumentUrl,
)

document_url_temporal_agent = TemporalAgent(document_url_agent, activity_config=BASE_ACTIVITY_CONFIG)


@workflow.defn
class DocumentUrlAgentWorkflow:
    @workflow.run
    async def run(self, prompt: str) -> DocumentUrl:
        result = await document_url_temporal_agent.run(prompt)
        return result.output


async def test_document_url_serialization_preserves_media_type(allow_model_requests: None, client: Client):
    """Test that `DocumentUrl` with custom `media_type` is preserved through Temporal serialization.

    This is a regression test for https://github.com/pydantic/pydantic-ai/issues/3949

Review comment (suggestion):
If we pass the `DocumentUrl` into `agent.run` instead of having a tool return it, it'll be easier to test this with `TestModel`.

Code after:
        task_queue=TASK_QUEUE,
        workflows=[ImageAgentWorkflow],
        plugins=[AgentPlugin(image_temporal_agent)],
    ):
        with workflow_raises(
            UserError,
            snapshot('Image output is not supported with Temporal because of the 2MB payload size limit.'),
        ):
            await client.execute_workflow(
                ImageAgentWorkflow.run,
                args=['Generate an image of an axolotl.'],
                id=ImageAgentWorkflow.__name__,
                task_queue=TASK_QUEUE,
            )


# ============================================================================
# DocumentUrl Serialization Test - Verifies that DocumentUrl with custom
# media_type is properly serialized through Temporal activities
# ============================================================================

document_url_agent = Agent(
    TestModel(custom_output_args={'url': 'https://example.com/doc/12345', 'media_type': 'application/pdf'}),
    name='document_url_agent',
    output_type=DocumentUrl,
)

document_url_temporal_agent = TemporalAgent(document_url_agent, activity_config=BASE_ACTIVITY_CONFIG)


@workflow.defn
class DocumentUrlAgentWorkflow:
    @workflow.run
    async def run(self, prompt: str) -> DocumentUrl:
        result = await document_url_temporal_agent.run(prompt)
        return result.output


async def test_document_url_serialization_preserves_media_type(allow_model_requests: None, client: Client):
    """Test that `DocumentUrl` with custom `media_type` is preserved through Temporal serialization.

    This is a regression test for https://github.com/pydantic/pydantic-ai/issues/3949
    where `DocumentUrl.media_type` (a computed field) was lost during Temporal activity
    serialization because the backing field `_media_type` was excluded from serialization.
    """
    async with Worker(
        client,
        task_queue=TASK_QUEUE,
        workflows=[DocumentUrlAgentWorkflow],
        plugins=[AgentPlugin(document_url_temporal_agent)],
    ):