chore: import upstream snapshot with attribution

tests/algorithm/__init__.py

@@ -0,0 +1 @@
+# Copyright (c) Microsoft. All rights reserved.

tests/algorithm/test_apo.py

@@ -0,0 +1,942 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+# pyright: reportPrivateUsage=false
+
+from typing import Any, Dict, Iterator, List, Literal, Optional, Sequence, Tuple, cast
+from unittest.mock import AsyncMock, Mock
+
+import pytest
+from openai import AsyncOpenAI
+
+import agentlightning.algorithm.apo.apo as apo_module
+from agentlightning.adapter import TraceAdapter
+from agentlightning.adapter.messages import TraceToMessages
+from agentlightning.algorithm.apo.apo import APO, RolloutResultForAPO, VersionedPromptTemplate, batch_iter_over_dataset
+from agentlightning.semconv import AGL_ANNOTATION
+from agentlightning.types import (
+    Dataset,
+    NamedResources,
+)
+from agentlightning.types import OtelResource as SpanResource
+from agentlightning.types import (
+    PromptTemplate,
+    Rollout,
+    Span,
+    SpanContext,
+    TraceStatus,
+)
+
+
+class DummyTraceMessagesAdapter(TraceToMessages):
+    def __init__(self) -> None:
+        super().__init__()
+        self.seen_spans: Sequence[Span] | None = None
+
+    def adapt(self, source: Sequence[Span], /) -> List[Dict[str, Any]]:  # type: ignore[override]
+        self.seen_spans = list(source)
+        return [dict(payload="converted")]
+
+
+class WrongAdapter(TraceAdapter[List[int]]):
+    def adapt(self, source: Sequence[Span], /) -> List[int]:
+        return [len(source)]
+
+
+class DummyStore:
+    def __init__(self) -> None:
+        self.update_resources_calls: List[Tuple[str, NamedResources]] = []
+        self.enqueue_calls: List[Dict[str, Any]] = []
+        self.wait_calls: List[Dict[str, Any]] = []
+        self.wait_results_queue: List[List[Rollout]] = []
+        self.query_spans_map: Dict[str, List[Span]] = {}
+        self._counter = 0
+
+    async def update_resources(self, resources_id: str, resources: NamedResources) -> Mock:
+        self.update_resources_calls.append((resources_id, resources))
+        update_mock = Mock()
+        update_mock.resources_id = resources_id
+        return update_mock
+
+    async def enqueue_rollout(
+        self,
+        *,
+        input: Dict[str, Any],
+        mode: str,
+        resources_id: Optional[str] = None,
+    ) -> Mock:
+        rollout_id = f"rollout-{self._counter}"
+        self._counter += 1
+        self.enqueue_calls.append(
+            {"rollout_id": rollout_id, "input": input, "mode": mode, "resources_id": resources_id}
+        )
+        result = Mock()
+        result.rollout_id = rollout_id
+        return result
+
+    async def wait_for_rollouts(self, rollout_ids: Sequence[str], timeout: float) -> List[Rollout]:
+        self.wait_calls.append({"rollout_ids": tuple(rollout_ids), "timeout": timeout})
+        if self.wait_results_queue:
+            return self.wait_results_queue.pop(0)
+        return []
+
+    async def query_spans(
+        self,
+        rollout_id: str,
+        attempt_id: str | Literal["latest"] | None = None,
+        **_: Any,
+    ) -> List[Span]:
+        return list(self.query_spans_map.get(rollout_id, []))
+
+
+def make_completion(content: str | None) -> Mock:
+    """Create a mock OpenAI completion response."""
+    message_mock = Mock()
+    message_mock.content = content
+    choice_mock = Mock()
+    choice_mock.message = message_mock
+    completion_mock = Mock()
+    completion_mock.choices = [choice_mock]
+    return completion_mock
+
+
+def make_openai_client(create_mock: AsyncMock) -> Mock:
+    """Create a mock AsyncOpenAI client with the given create method."""
+    client = Mock(spec=AsyncOpenAI)
+    completions = Mock()
+    completions.create = create_mock
+    chat = Mock()
+    chat.completions = completions
+    client.chat = chat
+    return client
+
+
+def make_reward_span(rollout_id: str, attempt_id: str, reward: float, sequence_id: int) -> Span:
+    hex_id = f"{sequence_id:032x}"
+    span_hex = f"{sequence_id:016x}"
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=hex_id,
+        span_id=span_hex,
+        parent_id=None,
+        name=AGL_ANNOTATION,
+        status=TraceStatus(status_code="OK"),
+        attributes={"reward": reward},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id=hex_id, span_id=span_hex, is_remote=False, trace_state={}),
+        parent=None,
+        resource=SpanResource(attributes={}, schema_url=""),
+    )
+
+
+def test_batch_iter_over_dataset_returns_full_dataset(monkeypatch: pytest.MonkeyPatch) -> None:
+    dataset = [{"id": idx} for idx in range(3)]
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+
+    iterator = batch_iter_over_dataset(cast(Dataset[Any], dataset), batch_size=5)
+
+    first_batch = next(iterator)
+    second_batch = next(iterator)
+
+    assert len(first_batch) == len(dataset)
+    assert len(second_batch) == len(dataset)
+    assert {item["id"] for item in first_batch} == {0, 1, 2}
+
+
+def test_batch_iter_over_dataset_cycles_batches(monkeypatch: pytest.MonkeyPatch) -> None:
+    dataset = [{"id": idx} for idx in range(4)]
+
+    def fake_shuffle(seq: List[int]) -> None:
+        seq.reverse()
+
+    monkeypatch.setattr(apo_module.random, "shuffle", fake_shuffle)
+
+    iterator = batch_iter_over_dataset(cast(Dataset[Any], dataset), batch_size=2)
+
+    batch_one = next(iterator)
+    batch_two = next(iterator)
+    batch_three = next(iterator)
+
+    assert len(batch_one) == 2
+    assert len(batch_two) == 2
+    assert {item["id"] for item in batch_three}.issubset({item["id"] for item in batch_one + batch_two})  # type: ignore
+
+
+def test_apo_init_sets_configuration() -> None:
+    client = Mock(spec=AsyncOpenAI)
+
+    apo = APO[Any](
+        client,
+        gradient_model="g-model",
+        apply_edit_model="a-model",
+        diversity_temperature=0.7,
+        gradient_batch_size=3,
+        val_batch_size=5,
+        beam_width=2,
+        branch_factor=3,
+        beam_rounds=4,
+        rollout_batch_timeout=42.0,
+        run_initial_validation=False,
+    )
+
+    assert apo.async_openai_client is client
+    assert apo.gradient_model == "g-model"
+    assert apo.apply_edit_model == "a-model"
+    assert apo.diversity_temperature == 0.7
+    assert apo.gradient_batch_size == 3
+    assert apo.val_batch_size == 5
+    assert apo.beam_width == 2
+    assert apo.branch_factor == 3
+    assert apo.beam_rounds == 4
+    assert apo.rollout_batch_timeout == 42.0
+    assert apo.run_initial_validation is False
+    assert apo._history_best_prompt is None
+    assert apo._history_best_score == float("-inf")
+
+
+def test_get_seed_prompt_template_returns_prompt() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    prompt = PromptTemplate(template="Seed: {x}", engine="f-string")
+    resources: NamedResources = {
+        "seed": prompt,
+        "other": PromptTemplate(template="Other", engine="f-string"),
+    }
+    apo.set_initial_resources(resources)
+
+    resource_name, seed_prompt = apo.get_seed_prompt_template()
+
+    assert resource_name == "seed"
+    assert seed_prompt is prompt
+
+
+def test_get_seed_prompt_template_requires_resources() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+
+    with pytest.raises(ValueError):
+        apo.get_seed_prompt_template()
+
+
+def test_get_seed_prompt_template_requires_prompt_resource() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    apo.set_initial_resources({})
+
+    with pytest.raises(ValueError):
+        apo.get_seed_prompt_template()
+
+
+def test_get_adapter_returns_trace_messages_adapter() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_adapter(adapter)
+
+    assert apo.get_adapter() is adapter
+
+
+def test_get_adapter_requires_trace_messages_adapter() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+    apo.set_adapter(WrongAdapter())
+
+    with pytest.raises(ValueError):
+        apo.get_adapter()
+
+
+def test_get_best_prompt_requires_history() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+
+    with pytest.raises(ValueError):
+        apo.get_best_prompt()
+
+
+def test_get_best_prompt_returns_prompt() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    prompt = PromptTemplate(template="Best", engine="f-string")
+    apo._history_best_prompt = prompt
+
+    assert apo.get_best_prompt() is prompt
+
+
+@pytest.mark.asyncio
+async def test_compute_textual_gradient_samples_batch(monkeypatch: pytest.MonkeyPatch) -> None:
+    create_mock = AsyncMock(return_value=make_completion("critique"))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_model="test-gradient-model", gradient_batch_size=2, diversity_temperature=0.8)
+    versioned_prompt = apo._create_versioned_prompt(PromptTemplate(template="prompt", engine="f-string"))
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=float(i), spans=[], messages=[]) for i in range(3)
+    ]
+
+    sample_mock = Mock(return_value=rollouts[:2])
+    monkeypatch.setattr(apo_module.random, "sample", sample_mock)
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    result = await apo.compute_textual_gradient(versioned_prompt, rollouts)
+
+    assert result == "critique"
+    sample_mock.assert_called_once_with(rollouts, 2)
+    # Verify OpenAI call was made with correct parameters
+    create_mock.assert_awaited_once()
+    call_kwargs = create_mock.await_args.kwargs  # type: ignore
+    assert call_kwargs["model"] == "test-gradient-model"
+    assert call_kwargs["temperature"] == 0.8
+    assert len(call_kwargs["messages"]) == 1
+    assert call_kwargs["messages"][0]["role"] == "user"
+    assert call_kwargs["messages"][0]["content"].startswith("You optimize a prompt template.")
+
+
+@pytest.mark.asyncio
+async def test_compute_textual_gradient_uses_all_rollouts_when_insufficient(monkeypatch: pytest.MonkeyPatch) -> None:
+    create_mock = AsyncMock(return_value=make_completion("critique"))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_batch_size=3)
+    versioned_prompt = apo._create_versioned_prompt(PromptTemplate(template="prompt", engine="f-string"))
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    sample_mock = Mock(side_effect=AssertionError("sample should not be called"))
+    monkeypatch.setattr(apo_module.random, "sample", sample_mock)
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    result = await apo.compute_textual_gradient(versioned_prompt, rollouts)
+
+    assert result == "critique"
+
+
+@pytest.mark.asyncio
+async def test_textual_gradient_and_apply_edit_returns_new_prompt(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Use two separate mocks for gradient and edit calls
+    gradient_mock = AsyncMock(return_value=make_completion("critique text"))
+    edit_mock = AsyncMock(return_value=make_completion("new prompt"))
+
+    call_count = 0
+
+    async def create_side_effect(*args: Any, **kwargs: Any) -> Mock:
+        nonlocal call_count
+        call_count += 1
+        return gradient_mock.return_value if call_count == 1 else edit_mock.return_value
+
+    create_mock = AsyncMock(side_effect=create_side_effect)
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client, gradient_model="grad-model", apply_edit_model="edit-model", diversity_temperature=0.9)
+
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+
+    poml_calls: List[Dict[str, Any]] = []
+
+    def poml_side_effect(template: Any, context: Dict[str, Any], format: str) -> Dict[str, Any]:
+        poml_calls.append({"template": template, "context": context, "format": format})
+        return {"messages": [{"role": "user", "content": "msg"}]}
+
+    monkeypatch.setattr(apo_module.poml, "poml", poml_side_effect)
+
+    versioned_prompt = apo._create_versioned_prompt(PromptTemplate(template="old prompt", engine="f-string"))
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    result = await apo.textual_gradient_and_apply_edit(versioned_prompt, rollouts)
+
+    assert result == "new prompt"
+    assert create_mock.await_count == 2
+
+    # Verify gradient computation call
+    first_call = create_mock.await_args_list[0].kwargs
+    assert first_call["model"] == "grad-model"
+    assert first_call["temperature"] == 0.9
+
+    # Verify edit application call
+    second_call = create_mock.await_args_list[1].kwargs
+    assert second_call["model"] == "edit-model"
+    assert second_call["temperature"] == 0.9
+
+    # Verify critique was passed to edit context
+    assert len(poml_calls) == 2
+    assert poml_calls[1]["context"]["critique"] == "critique text"
+    assert poml_calls[1]["context"]["prompt_template"] == "old prompt"
+
+
+@pytest.mark.asyncio
+async def test_textual_gradient_and_apply_edit_returns_original_if_no_critique(monkeypatch: pytest.MonkeyPatch) -> None:
+    # Mock OpenAI to return None content
+    create_mock = AsyncMock(return_value=make_completion(None))
+    client = make_openai_client(create_mock)
+    apo = APO[Any](client)
+
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+
+    versioned_prompt = apo._create_versioned_prompt(PromptTemplate(template="old prompt", engine="f-string"))
+    rollouts: List[RolloutResultForAPO] = [
+        RolloutResultForAPO(status="succeeded", final_reward=1.0, spans=[], messages=[])
+    ]
+
+    result = await apo.textual_gradient_and_apply_edit(versioned_prompt, rollouts)
+
+    # Should return original prompt when gradient computation fails
+    assert result == "old prompt"
+    # Verify gradient computation was attempted
+    create_mock.assert_awaited_once()
+
+
+@pytest.mark.asyncio
+async def test_get_rollout_results_adapts_spans() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    rollout = Rollout(
+        rollout_id="r-1",
+        input={"task": "value"},
+        start_time=0.0,
+        status="succeeded",
+        mode="train",
+    )
+    span1 = make_reward_span("r-1", "attempt", 1.0, sequence_id=1)
+    span2 = make_reward_span("r-1", "attempt", 2.0, sequence_id=2)
+    store.query_spans_map["r-1"] = [span1, span2]
+
+    results = await apo.get_rollout_results([rollout])
+
+    assert len(results) == 1
+    # Verify final reward is correctly extracted
+    assert results[0]["final_reward"] == 2.0
+    # Verify status is correctly mapped
+    assert results[0]["status"] == "succeeded"
+    # Verify adapter was called with correct spans
+    assert adapter.seen_spans is not None
+    assert len(adapter.seen_spans) == 2
+    assert adapter.seen_spans[0] == span1
+    assert adapter.seen_spans[1] == span2
+    # Verify messages were converted
+    assert results[0]["messages"] == [{"payload": "converted"}]
+    # Verify spans were serialized
+    assert len(results[0]["spans"]) == 2
+    assert results[0]["spans"][0]["rollout_id"] == "r-1"
+    assert results[0]["spans"][0]["name"] == AGL_ANNOTATION
+    assert results[0]["spans"][0]["attributes"]["reward"] == 1.0
+    assert results[0]["spans"][1]["attributes"]["reward"] == 2.0
+
+
+@pytest.mark.asyncio
+async def test_evaluate_prompt_on_batch_runs_rollouts() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, rollout_batch_timeout=100.0)
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    dataset = [{"task": 1}, {"task": 2}]
+
+    # Set up spans for rollouts
+    store.query_spans_map["rollout-0"] = [make_reward_span("rollout-0", "attempt", 1.0, sequence_id=1)]
+    store.query_spans_map["rollout-1"] = [make_reward_span("rollout-1", "attempt", 0.0, sequence_id=1)]
+
+    store.wait_results_queue.append(
+        [
+            Rollout(
+                rollout_id="rollout-0",
+                input=dataset[0],
+                start_time=0.0,
+                status="succeeded",
+                mode="train",
+            ),
+            Rollout(
+                rollout_id="rollout-1",
+                input=dataset[1],
+                start_time=0.0,
+                status="failed",
+                mode="train",
+            ),
+        ]
+    )
+
+    prompt_template = PromptTemplate(template="test prompt", engine="f-string")
+    versioned_prompt = apo._create_versioned_prompt(prompt_template)
+
+    rollout_results, average = await apo.evaluate_prompt_on_batch(versioned_prompt, "seed", dataset, mode="train")
+
+    # Verify results
+    assert len(rollout_results) == 2
+    assert rollout_results[0]["final_reward"] == 1.0
+    assert rollout_results[0]["status"] == "succeeded"
+    assert rollout_results[1]["final_reward"] == 0.0
+    assert rollout_results[1]["status"] == "failed"
+    assert average == pytest.approx(0.5)  # type: ignore
+
+    # Verify resource was added with correct prompt
+    assert len(store.update_resources_calls) == 1
+    resources_id, resources_payload = store.update_resources_calls[0]
+    assert resources_id == versioned_prompt.version
+    assert "seed" in resources_payload
+    added_resource = resources_payload["seed"]
+    assert isinstance(added_resource, PromptTemplate)
+    assert added_resource.template == "test prompt"
+    assert added_resource.engine == "f-string"
+
+    # Verify enqueue was called correctly
+    assert len(store.enqueue_calls) == 2
+    assert store.enqueue_calls[0]["input"] == dataset[0]
+    assert store.enqueue_calls[0]["mode"] == "train"
+    assert store.enqueue_calls[0]["resources_id"] == versioned_prompt.version
+    assert store.enqueue_calls[1]["input"] == dataset[1]
+    assert store.enqueue_calls[1]["mode"] == "train"
+    assert store.enqueue_calls[1]["resources_id"] == versioned_prompt.version
+
+    # Verify wait was called with correct rollout IDs
+    assert len(store.wait_calls) == 1
+    assert set(store.wait_calls[0]["rollout_ids"]) == {"rollout-0", "rollout-1"}
+    assert store.wait_calls[0]["timeout"] == 0.0
+
+
+def test_initialize_beam_sets_history(monkeypatch: pytest.MonkeyPatch) -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, gradient_batch_size=2, val_batch_size=1)
+    prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": prompt})
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+
+    train_dataset: Sequence[Dict[str, str]] = [{"x": "1"}, {"x": "2"}]
+    val_dataset: Sequence[Dict[str, str]] = [{"y": "value"}]
+
+    resource_name, seed_prompt, grad_iter, val_iter = apo._initialize_beam(train_dataset, val_dataset)  # type: ignore
+
+    assert resource_name == "seed"
+    assert seed_prompt is prompt
+    assert apo._history_best_prompt is prompt
+    assert apo._history_best_score == float("-inf")
+    assert len(next(grad_iter)) == len(train_dataset)
+    assert len(next(val_iter)) == len(val_dataset)
+
+
+def test_initialize_beam_requires_train_dataset() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    apo.set_initial_resources({"seed": PromptTemplate(template="Seed", engine="f-string")})
+
+    with pytest.raises(ValueError):
+        apo._initialize_beam(None, [])  # type: ignore
+
+
+def test_initialize_beam_requires_val_dataset() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    apo.set_initial_resources({"seed": PromptTemplate(template="Seed", engine="f-string")})
+
+    with pytest.raises(ValueError):
+        apo._initialize_beam([], None)  # type: ignore
+
+
+def test_sample_parent_prompts_replicates_when_beam_too_small(monkeypatch: pytest.MonkeyPatch) -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=3)
+    beam_prompt = apo._create_versioned_prompt(PromptTemplate(template="Seed", engine="f-string"))
+    beam = [beam_prompt]
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: (_ for _ in ()).throw(AssertionError()))  # type: ignore
+
+    sampled = apo._sample_parent_prompts(beam, round_num=0)
+
+    assert len(sampled) == apo.beam_width
+    assert all(index == 0 and prompt is beam_prompt for index, prompt in sampled)
+
+
+def test_sample_parent_prompts_uses_random_sample(monkeypatch: pytest.MonkeyPatch) -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=2)
+    prompt_a = apo._create_versioned_prompt(PromptTemplate(template="A", engine="f-string"))
+    prompt_b = apo._create_versioned_prompt(PromptTemplate(template="B", engine="f-string"))
+    prompt_c = apo._create_versioned_prompt(PromptTemplate(template="C", engine="f-string"))
+
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: [0, 2])  # type: ignore
+
+    sampled = apo._sample_parent_prompts([prompt_a, prompt_b, prompt_c], round_num=1)
+
+    assert sampled == [(0, prompt_a), (2, prompt_c)]
+
+
+@pytest.mark.asyncio
+async def test_generate_candidate_prompts_creates_branch_factor_children() -> None:
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, branch_factor=2)
+    store = DummyStore()
+    adapter = DummyTraceMessagesAdapter()
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    parent_prompt = apo._create_versioned_prompt(PromptTemplate(template="Seed", engine="f-string"))
+    grad_batches: Iterator[Sequence[Dict[str, Any]]] = iter(
+        [
+            [{"task": "a"}],
+            [{"task": "b"}],
+        ]
+    )
+
+    # Set up rollouts to complete immediately
+    store.query_spans_map["rollout-0"] = [make_reward_span("rollout-0", "attempt", 0.5, sequence_id=1)]
+    store.query_spans_map["rollout-1"] = [make_reward_span("rollout-1", "attempt", 0.6, sequence_id=1)]
+    store.wait_results_queue.extend(
+        [
+            [Rollout(rollout_id="rollout-0", input={"task": "a"}, start_time=0.0, status="succeeded", mode="train")],
+            [Rollout(rollout_id="rollout-1", input={"task": "b"}, start_time=0.0, status="succeeded", mode="train")],
+        ]
+    )
+
+    counter = 0
+
+    async def edit_side_effect(
+        current_prompt: VersionedPromptTemplate,
+        rollout: List[RolloutResultForAPO],
+        **_: Any,
+    ) -> str:
+        nonlocal counter
+        counter += 1
+        return f"{current_prompt.prompt_template.template}-{counter}"
+
+    apo.textual_gradient_and_apply_edit = AsyncMock(side_effect=edit_side_effect)
+
+    candidates = await apo._generate_candidate_prompts([(0, parent_prompt)], "seed", grad_batches, round_num=0)
+
+    # Verify correct number of candidates generated
+    assert len(candidates) == apo.branch_factor
+    assert {candidate.prompt_template.template for candidate in candidates} == {"Seed-1", "Seed-2"}
+    assert all(candidate.prompt_template.engine == "f-string" for candidate in candidates)
+
+    # Verify evaluate_prompt_on_batch was called for each candidate generation
+    assert len(store.enqueue_calls) == 2
+    assert store.enqueue_calls[0]["input"] == {"task": "a"}
+    assert store.enqueue_calls[1]["input"] == {"task": "b"}
+    assert all(call["mode"] == "train" for call in store.enqueue_calls)
+    assert all(call["resources_id"] == parent_prompt.version for call in store.enqueue_calls)
+
+    # Verify textual_gradient_and_apply_edit was called correct number of times
+    assert apo.textual_gradient_and_apply_edit.await_count == 2
+
+
+@pytest.mark.asyncio
+async def test_generate_candidate_prompts_skips_failed_generations() -> None:
+    """Test that None returns from textual_gradient_and_apply_edit are skipped."""
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, branch_factor=3)
+    store = DummyStore()
+    # Keep strong reference to prevent garbage collection since APO uses weakref
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    parent_prompt = apo._create_versioned_prompt(PromptTemplate(template="Seed", engine="f-string"))
+    grad_batches: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": f"t{i}"}] for i in range(3)])
+
+    # Set up rollouts
+    for i in range(3):
+        store.query_spans_map[f"rollout-{i}"] = [make_reward_span(f"rollout-{i}", "attempt", 0.5, sequence_id=1)]
+        store.wait_results_queue.append(
+            [
+                Rollout(
+                    rollout_id=f"rollout-{i}", input={"task": f"t{i}"}, start_time=0.0, status="succeeded", mode="train"
+                )
+            ]
+        )
+
+    # Mock to return None for second call, valid prompts for others
+    call_count = 0
+
+    async def edit_side_effect(
+        current_prompt: VersionedPromptTemplate,
+        rollout: List[RolloutResultForAPO],
+        **_: Any,
+    ) -> Optional[str]:
+        nonlocal call_count
+        call_count += 1
+        if call_count == 2:
+            return None  # Simulate failure
+        return f"{current_prompt.prompt_template.template}-{call_count}"
+
+    apo.textual_gradient_and_apply_edit = AsyncMock(side_effect=edit_side_effect)
+
+    candidates = await apo._generate_candidate_prompts([(0, parent_prompt)], "seed", grad_batches, round_num=0)
+
+    # Should only have 2 candidates (one failed)
+    assert len(candidates) == 2
+    assert {candidate.prompt_template.template for candidate in candidates} == {"Seed-1", "Seed-3"}
+    # Verify all three attempts were made
+    assert apo.textual_gradient_and_apply_edit.await_count == 3
+
+
+@pytest.mark.asyncio
+async def test_evaluate_and_select_beam_sorts_by_score() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI), beam_width=2)
+    candidates = [
+        apo._create_versioned_prompt(PromptTemplate(template="A", engine="f-string")),
+        apo._create_versioned_prompt(PromptTemplate(template="B", engine="f-string")),
+        apo._create_versioned_prompt(PromptTemplate(template="C", engine="f-string")),
+    ]
+    scores = {"A": 1.0, "B": 0.2, "C": 2.0}
+
+    async def evaluate(
+        prompt: VersionedPromptTemplate,
+        resource_name: str,
+        dataset: Sequence[Dict[str, Any]],
+        mode: str,
+        **_: Any,
+    ) -> Any:
+        return [], scores[prompt.prompt_template.template]
+
+    apo.evaluate_prompt_on_batch = AsyncMock(side_effect=evaluate)  # type: ignore[assignment]
+
+    val_iterator: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": "val"}]])
+
+    selected = await apo._evaluate_and_select_beam(candidates, "seed", val_iterator, round_num=0)
+
+    assert [prompt.prompt_template.template for prompt in selected] == ["C", "A"]
+
+
+@pytest.mark.asyncio
+async def test_evaluate_and_select_beam_raises_on_empty_candidates() -> None:
+    """Test that ValueError is raised when no candidates remain after evaluation."""
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client, beam_width=2)
+    # Empty candidate list
+    candidates: List[VersionedPromptTemplate] = []
+
+    val_iterator: Iterator[Sequence[Dict[str, Any]]] = iter([[{"task": "val"}]])
+
+    with pytest.raises(ValueError, match="No beam candidates any more"):
+        await apo._evaluate_and_select_beam(candidates, "seed", val_iterator, round_num=0)
+
+
+@pytest.mark.asyncio
+async def test_update_best_prompt_updates_history() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    old_versioned = apo._create_versioned_prompt(PromptTemplate(template="Old", engine="f-string"))
+    new_versioned = apo._create_versioned_prompt(PromptTemplate(template="New", engine="f-string"))
+    apo._history_best_prompt = old_versioned.prompt_template
+    apo._history_best_score = 0.5
+    apo._history_best_version = old_versioned.version
+    apo.evaluate_prompt_on_batch = AsyncMock(return_value=([], 1.2))  # type: ignore[assignment]
+
+    await apo._update_best_prompt([new_versioned], "seed", [{"task": "val"}], round_num=0)  # type: ignore
+
+    assert apo._history_best_prompt is new_versioned.prompt_template
+    assert apo._history_best_score == pytest.approx(1.2)  # type: ignore
+    assert apo._history_best_version == new_versioned.version
+
+
+@pytest.mark.asyncio
+async def test_update_best_prompt_keeps_history_when_not_improved() -> None:
+    apo = APO[Any](Mock(spec=AsyncOpenAI))
+    old_versioned = apo._create_versioned_prompt(PromptTemplate(template="Old", engine="f-string"))
+    new_versioned = apo._create_versioned_prompt(PromptTemplate(template="New", engine="f-string"))
+    apo._history_best_prompt = old_versioned.prompt_template
+    apo._history_best_score = 2.0
+    apo._history_best_version = old_versioned.version
+    apo.evaluate_prompt_on_batch = AsyncMock(return_value=([], 1.5))  # type: ignore[assignment]
+
+    await apo._update_best_prompt([new_versioned], "seed", [{"task": "val"}], round_num=0)  # type: ignore
+
+    assert apo._history_best_prompt is old_versioned.prompt_template
+    assert apo._history_best_score == pytest.approx(2.0)  # type: ignore
+    assert apo._history_best_version == old_versioned.version
+
+
+def test_apo_init_defaults_run_initial_validation_to_true() -> None:
+    """Test that run_initial_validation defaults to True when not specified."""
+    client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](client)
+
+    assert apo.run_initial_validation is True
+
+
+@pytest.mark.asyncio
+async def test_run_performs_initial_validation_when_enabled(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that initial validation runs on seed prompt when run_initial_validation=True."""
+    async_client = Mock(spec=AsyncOpenAI)
+    apo = APO[Any](
+        async_client,
+        gradient_batch_size=1,
+        val_batch_size=1,
+        beam_width=1,
+        branch_factor=1,
+        beam_rounds=0,  # No optimization rounds, just initial validation
+        run_initial_validation=True,
+    )
+    seed_prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": seed_prompt})
+
+    store = DummyStore()
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    # Set up initial validation rollout
+    store.query_spans_map["rollout-0"] = [make_reward_span("rollout-0", "attempt", 0.75, sequence_id=1)]
+    store.wait_results_queue.append(
+        [
+            Rollout(
+                rollout_id="rollout-0",
+                input={"task": "val"},
+                start_time=0.0,
+                status="succeeded",
+                mode="val",
+            )
+        ]
+    )
+
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+
+    val_dataset = [{"task": "val"}]
+    await apo.run(train_dataset=[{"task": "train"}], val_dataset=val_dataset)  # type: ignore
+
+    # Verify initial validation was performed
+    assert apo._history_best_prompt is seed_prompt
+    assert apo._history_best_score == pytest.approx(0.75)  # type: ignore
+
+    # Verify a validation rollout was enqueued for initial validation
+    val_calls = [c for c in store.enqueue_calls if c["mode"] == "val"]
+    assert len(val_calls) == 1
+
+
+@pytest.mark.asyncio
+async def test_run_skips_initial_validation_when_disabled(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Test that initial validation is skipped when run_initial_validation=False."""
+    create_mock = AsyncMock(side_effect=[make_completion("critique text"), make_completion("improved prompt")])
+    async_client = make_openai_client(create_mock)
+
+    apo = APO[Any](
+        async_client,
+        gradient_batch_size=1,
+        val_batch_size=1,
+        beam_width=1,
+        branch_factor=1,
+        beam_rounds=1,
+        run_initial_validation=False,  # Disable initial validation
+    )
+    seed_prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": seed_prompt})
+
+    store = DummyStore()
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    # Set up spans for rollouts (train + val for candidates + final val)
+    rollout_rewards = [0.4, 0.5, 0.6, 1.1]
+    for i, reward in enumerate(rollout_rewards):
+        store.query_spans_map[f"rollout-{i}"] = [make_reward_span(f"rollout-{i}", "attempt", reward, sequence_id=1)]
+        store.wait_results_queue.append(
+            [
+                Rollout(
+                    rollout_id=f"rollout-{i}",
+                    input={"task": f"data-{i}"},
+                    start_time=0.0,
+                    status="succeeded",
+                    mode="train" if i == 0 else "val",
+                )
+            ]
+        )
+
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    train_dataset = [{"task": "train"}]
+    val_dataset = [{"task": "val"}]
+
+    await apo.run(train_dataset=train_dataset, val_dataset=val_dataset)  # type: ignore
+
+    # Verify best prompt was updated through normal optimization (not initial validation)
+    best_prompt = apo.get_best_prompt()
+    assert best_prompt.template == "improved prompt"
+
+    # Count validation rollouts - should NOT include initial validation
+    # Only candidate evaluation + final best prompt evaluation
+    val_calls = [c for c in store.enqueue_calls if c["mode"] == "val"]
+    # With run_initial_validation=False, we expect: 2 val calls (seed+new candidate) + 1 final val = 3 total
+    assert len(val_calls) == 3
+
+
+@pytest.mark.asyncio
+async def test_run_updates_best_prompt_with_real_openai_client(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Integration test for the full run method with minimal mocking."""
+    create_mock = AsyncMock(side_effect=[make_completion("critique text"), make_completion("improved prompt")])
+    async_client = make_openai_client(create_mock)
+
+    apo = APO[Any](
+        async_client,
+        gradient_batch_size=1,
+        val_batch_size=1,
+        beam_width=1,
+        branch_factor=1,
+        beam_rounds=1,
+        run_initial_validation=False,  # Skip initial validation for this test
+    )
+    seed_prompt = PromptTemplate(template="Seed", engine="f-string")
+    apo.set_initial_resources({"seed": seed_prompt})
+
+    store = DummyStore()
+    # Keep strong reference to prevent garbage collection since APO uses weakref
+    apo._test_adapter = adapter = DummyTraceMessagesAdapter()  # type: ignore
+    apo.set_store(store)  # type: ignore
+    apo.set_adapter(adapter)
+
+    # Set up spans for all expected rollouts
+    # For 1 round with beam_width=1, branch_factor=1, run_initial_validation=False, we expect:
+    # 1. Training rollout for gradient computation
+    # 2. Validation rollouts for candidate evaluation (seed + new candidate = 2)
+    # 3. Final validation rollout on full dataset for best prompt
+    rollout_rewards = [0.4, 0.5, 0.6, 1.1]
+    for i, reward in enumerate(rollout_rewards):
+        store.query_spans_map[f"rollout-{i}"] = [make_reward_span(f"rollout-{i}", "attempt", reward, sequence_id=1)]
+        store.wait_results_queue.append(
+            [
+                Rollout(
+                    rollout_id=f"rollout-{i}",
+                    input={"task": f"data-{i}"},
+                    start_time=0.0,
+                    status="succeeded",
+                    mode="train" if i == 0 else "val",
+                )
+            ]
+        )
+
+    monkeypatch.setattr(apo_module.random, "shuffle", lambda seq: None)  # type: ignore
+    monkeypatch.setattr(apo_module.random, "sample", lambda population, k: list(population)[:k])  # type: ignore
+    monkeypatch.setattr(apo_module.random, "choice", lambda seq: seq[0])  # type: ignore
+
+    train_dataset = [{"task": "train"}]
+    val_dataset = [{"task": "val"}]
+
+    await apo.run(train_dataset=train_dataset, val_dataset=val_dataset)  # type: ignore
+
+    # Verify best prompt was updated
+    best_prompt = apo.get_best_prompt()
+    assert best_prompt.template == "improved prompt"
+
+    # Verify OpenAI was called twice (gradient + edit)
+    assert create_mock.await_count == 2
+    gradient_call = create_mock.await_args_list[0]
+    assert gradient_call.kwargs["model"] == apo.gradient_model
+    edit_call = create_mock.await_args_list[1]
+    assert edit_call.kwargs["model"] == apo.apply_edit_model
+
+    # Verify resources were updated (seed prompt + new candidate prompts)
+    assert len(store.update_resources_calls) >= 2
+    assert all(isinstance(entry[0], str) for entry in store.update_resources_calls)
+
+    # Verify rollouts were enqueued (1 train + multiple val)
+    assert len(store.enqueue_calls) >= 3
+    train_calls = [c for c in store.enqueue_calls if c["mode"] == "train"]
+    val_calls = [c for c in store.enqueue_calls if c["mode"] == "val"]
+    assert len(train_calls) == 1
+    assert len(val_calls) >= 2
+
+    # Verify history was updated correctly
+    assert apo._history_best_prompt is not None
+    assert apo._history_best_score > 0

tests/algorithm/test_baseline.py

@@ -0,0 +1,249 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import logging
+from contextlib import suppress
+from dataclasses import dataclass
+from typing import Any, Dict, List, Sequence
+
+import pytest
+
+from agentlightning.adapter import TraceAdapter
+from agentlightning.algorithm import Baseline
+from agentlightning.store.memory import InMemoryLightningStore
+from agentlightning.types import (
+    LLM,
+    NamedResources,
+    OtelResource,
+    Span,
+    SpanContext,
+    TraceStatus,
+)
+
+LOGGER_NAME = "agentlightning.algorithm.fast"
+
+
+class _AdapterStub(TraceAdapter[Dict[str, Any]]):
+    def adapt(self, source: Sequence[Span], /) -> Dict[str, Any]:
+        return {
+            "count": len(source),
+            "attempt_ids": sorted({span.attempt_id for span in source}),
+        }
+
+
+@dataclass
+class _RolloutArtifacts:
+    rollout_id: str
+    attempt_id: str
+    attempt_sequence: int
+    span: Span
+
+
+def _make_resources() -> NamedResources:
+    return {
+        "main_llm": LLM(endpoint="http://localhost", model="test-model"),
+    }
+
+
+def _build_span(rollout_id: str, attempt_id: str, *, sequence_id: int, index: int) -> Span:
+    trace_hex = f"{index:032x}"
+    span_hex = f"{index:016x}"
+    # Minimal span that passes validation and keeps log output predictable.
+    return Span(
+        rollout_id=rollout_id,
+        attempt_id=attempt_id,
+        sequence_id=sequence_id,
+        trace_id=trace_hex,
+        span_id=span_hex,
+        parent_id=None,
+        name="test-span",
+        status=TraceStatus(status_code="OK"),
+        attributes={"stage": "collect"},
+        events=[],
+        links=[],
+        start_time=None,
+        end_time=None,
+        context=SpanContext(trace_id=trace_hex, span_id=span_hex, is_remote=False, trace_state={}),
+        parent=None,
+        resource=OtelResource(attributes={}, schema_url=""),
+    )
+
+
+async def _mock_runner(
+    *,
+    store: InMemoryLightningStore,
+    expected: int,
+    artifacts: List[_RolloutArtifacts],
+) -> None:
+    """Simulate a runner consuming rollouts, adding spans, and marking them complete."""
+    processed = 0
+    while processed < expected:
+        attempted = await store.dequeue_rollout()
+        if attempted is None:
+            await asyncio.sleep(0.001)
+            continue
+
+        attempt = attempted.attempt
+        rollout_id = attempted.rollout_id
+        await store.update_attempt(
+            rollout_id,
+            attempt.attempt_id,
+            status="running",
+            worker_id="runner-1",
+        )
+
+        span = _build_span(rollout_id, attempt.attempt_id, sequence_id=1, index=processed + 1)
+        await store.add_span(span)
+        await store.update_attempt(rollout_id, attempt.attempt_id, status="succeeded")
+        await store.update_rollout(rollout_id, status="succeeded")
+
+        artifacts.append(
+            _RolloutArtifacts(
+                rollout_id=rollout_id,
+                attempt_id=attempt.attempt_id,
+                attempt_sequence=attempt.sequence_id,
+                span=span,
+            )
+        )
+        processed += 1
+
+
+@pytest.mark.asyncio
+async def test_mock_algorithm_collects_rollout_logs(caplog: pytest.LogCaptureFixture) -> None:
+    store = InMemoryLightningStore()
+    await store.update_resources("default", _make_resources())
+    algorithm = Baseline(polling_interval=0.01, span_verbosity="key_values")
+    algorithm.set_store(store)
+    adapter = _AdapterStub()
+    algorithm.set_adapter(adapter)
+
+    caplog.set_level(logging.INFO, logger=LOGGER_NAME)
+
+    train_dataset = ["train-sample", "validation-sample"]
+    expected_rollouts = len(train_dataset)
+    artifacts: List[_RolloutArtifacts] = []
+
+    runner_task = asyncio.create_task(_mock_runner(store=store, expected=expected_rollouts, artifacts=artifacts))
+    try:
+        await algorithm.run(train_dataset=train_dataset)
+        await asyncio.wait_for(runner_task, timeout=2)
+    finally:
+        if not runner_task.done():
+            runner_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await runner_task
+
+    log_messages = [record.getMessage() for record in caplog.records if record.name == LOGGER_NAME]
+
+    # Ensure final status, attempt details, span details, and adapter output are logged per rollout.
+    for entry in artifacts:
+        attempt_summary = (
+            f"[Rollout {entry.rollout_id} | Attempt {entry.attempt_sequence}] "
+            f"ID: {entry.attempt_id}. Status: succeeded. Worker: runner-1"
+        )
+        assert attempt_summary in log_messages
+
+        span_prefix = (
+            f"[Rollout {entry.rollout_id} | Attempt {entry.attempt_id} | Span {entry.span.span_id}] "
+            f"#{entry.span.sequence_id} ({entry.span.name}) "
+        )
+        assert any(msg.startswith(span_prefix + "From") for msg in log_messages)
+        assert any(f"Attributes: {entry.span.attributes}" in msg for msg in log_messages)
+        assert any(
+            msg.startswith(f"[Rollout {entry.rollout_id}] Finished with status succeeded") for msg in log_messages
+        )
+        assert any(
+            msg.startswith(f"[Rollout {entry.rollout_id}] Adapted data: ")
+            and "'count': 1" in msg
+            and entry.attempt_id in msg
+            for msg in log_messages
+        )
+
+
+@pytest.mark.asyncio
+async def test_baseline_does_not_skip_samples_when_queue_full() -> None:
+    """Test that Baseline waits and retries when queue is full instead of skipping samples.
+
+    This is a regression test for a bug where samples would be skipped when the queue
+    exceeded max_queue_length. The fix wraps the queue check in a while loop to ensure
+    all samples are eventually processed.
+    """
+    store = InMemoryLightningStore()
+    await store.update_resources("default", _make_resources())
+
+    # Use a small max_queue_length and fast polling to test queue full behavior
+    algorithm = Baseline(polling_interval=0.01, max_queue_length=1)
+    algorithm.set_store(store)
+
+    # Create a dataset with 5 samples
+    train_dataset = [f"sample-{i}" for i in range(5)]
+    expected_rollouts = len(train_dataset)
+
+    # Track which samples were enqueued
+    enqueued_samples: List[Any] = []
+    artifacts: List[_RolloutArtifacts] = []
+
+    async def _slow_runner() -> None:
+        """A slow runner that creates backpressure by processing rollouts with delays."""
+        processed = 0
+        while processed < expected_rollouts:
+            attempted = await store.dequeue_rollout()
+            if attempted is None:
+                await asyncio.sleep(0.01)
+                continue
+
+            attempt = attempted.attempt
+            rollout_id = attempted.rollout_id
+            rollout = await store.get_rollout_by_id(rollout_id)
+
+            # Track the sample that was enqueued
+            if rollout:
+                enqueued_samples.append(rollout.input)
+
+            await store.update_attempt(
+                rollout_id,
+                attempt.attempt_id,
+                status="running",
+                worker_id="slow-runner",
+            )
+
+            # Add a delay to create backpressure and cause queue to fill up
+            await asyncio.sleep(0.05)
+
+            span = _build_span(rollout_id, attempt.attempt_id, sequence_id=1, index=processed + 1)
+            await store.add_span(span)
+            await store.update_attempt(rollout_id, attempt.attempt_id, status="succeeded")
+            await store.update_rollout(rollout_id, status="succeeded")
+
+            artifacts.append(
+                _RolloutArtifacts(
+                    rollout_id=rollout_id,
+                    attempt_id=attempt.attempt_id,
+                    attempt_sequence=attempt.sequence_id,
+                    span=span,
+                )
+            )
+            processed += 1
+
+    runner_task = asyncio.create_task(_slow_runner())
+    try:
+        await algorithm.run(train_dataset=train_dataset)
+        await asyncio.wait_for(runner_task, timeout=5)
+    finally:
+        if not runner_task.done():
+            runner_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await runner_task
+
+    # Verify that ALL samples were enqueued and processed (no samples skipped)
+    assert (
+        len(enqueued_samples) == expected_rollouts
+    ), f"Expected {expected_rollouts} samples to be enqueued, but got {len(enqueued_samples)}"
+    assert (
+        len(artifacts) == expected_rollouts
+    ), f"Expected {expected_rollouts} rollouts to be processed, but got {len(artifacts)}"
+
+    # Verify that the enqueued samples match the dataset (in order)
+    assert (
+        enqueued_samples == train_dataset
+    ), f"Enqueued samples {enqueued_samples} do not match dataset {train_dataset}"

tests/algorithm/test_decorator.py

@@ -0,0 +1,362 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Test that @algo decorator preserves function executability."""
+
+import inspect
+from typing import Any, Optional
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlightning.algorithm.decorator import FunctionalAlgorithm, algo
+from agentlightning.store.base import LightningStore
+from agentlightning.types import Dataset
+
+
+@algo
+def sample_algorithm_func(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+    """A test function with algorithm decorator."""
+    # Store the datasets in a way we can verify
+    sample_algorithm_func.last_train = train_dataset  # type: ignore
+    sample_algorithm_func.last_val = val_dataset  # type: ignore
+
+
+def test_algorithm_preserves_executability():
+    """Test that @algo decorated functions remain executable."""
+    test_train = ["train1", "train2"]
+    test_val = ["val1"]
+
+    # Function should be callable
+    assert callable(sample_algorithm_func)
+
+    # Function should execute with keyword arguments
+    sample_algorithm_func(train_dataset=test_train, val_dataset=test_val)
+
+    # Verify it was called with the right arguments
+    assert sample_algorithm_func.last_train == test_train  # type: ignore
+    assert sample_algorithm_func.last_val == test_val  # type: ignore
+
+
+def test_algorithm_preserves_metadata():
+    """Test that @algo preserves function metadata."""
+    # Function name should be preserved
+    assert sample_algorithm_func.__name__ == "sample_algorithm_func"  # type: ignore
+
+    # Docstring should be preserved
+    assert sample_algorithm_func.__doc__ == "A test function with algorithm decorator."
+
+
+def test_algorithm_returns_functional_algorithm_instance():
+    """Test that @algo returns a FunctionalAlgorithm instance."""
+    assert isinstance(sample_algorithm_func, FunctionalAlgorithm)
+
+    # Should have algorithm methods
+    assert hasattr(sample_algorithm_func, "run")
+    assert hasattr(sample_algorithm_func, "get_store")
+    assert hasattr(sample_algorithm_func, "set_trainer")
+
+
+def test_algorithm_preserves_signature():
+    """Test that @algo preserves function signature."""
+    sig = inspect.signature(sample_algorithm_func)
+    params = list(sig.parameters.keys())
+
+    # Should have the expected parameters
+    assert params == ["train_dataset", "val_dataset"]
+
+
+def test_algorithm_run_method():
+    """Test that the run method works correctly."""
+
+    @algo
+    def test_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """Test algorithm."""
+        test_algo.executed = True  # type: ignore
+        test_algo.train = train_dataset  # type: ignore
+        test_algo.val = val_dataset  # type: ignore
+
+    test_algo.executed = False  # type: ignore
+
+    train_data = ["item1", "item2"]
+    val_data = ["val1"]
+
+    # Call run method
+    test_algo.run(train_data, val_data)
+
+    # Verify execution
+    assert test_algo.executed  # type: ignore
+    assert test_algo.train == train_data  # type: ignore
+    assert test_algo.val == val_data  # type: ignore
+
+
+def test_algorithm_callable_shortcut():
+    """Test that calling the instance directly works."""
+
+    @algo
+    def test_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """Test algorithm."""
+        test_algo.called = True  # type: ignore
+
+    test_algo.called = False  # type: ignore
+
+    # Direct call should work with keyword arguments
+    test_algo(train_dataset=None, val_dataset=None)
+
+    assert test_algo.called  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_async_function_with_algorithm():
+    """Test that async functions work with @algo decorator."""
+
+    @algo
+    async def async_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """An async test function."""
+        async_algo.executed = True  # type: ignore
+        async_algo.train = train_dataset  # type: ignore
+
+    async_algo.executed = False  # type: ignore
+
+    # Should be callable
+    assert callable(async_algo)
+
+    # Should preserve async nature when called directly with keyword arguments
+    test_data = ["async-test"]
+    await async_algo(train_dataset=test_data, val_dataset=None)
+
+    assert async_algo.executed  # type: ignore
+    assert async_algo.train == test_data  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_async_algorithm_run_method():
+    """Test that async algorithms work with the run method."""
+
+    @algo
+    async def async_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """An async algorithm."""
+        async_algo.run_executed = True  # type: ignore
+        async_algo.run_train = train_dataset  # type: ignore
+        async_algo.run_val = val_dataset  # type: ignore
+
+    async_algo.run_executed = False  # type: ignore
+
+    train_data = ["async-train"]
+    val_data = ["async-val"]
+
+    # Run method should return an awaitable
+    assert async_algo.is_async()
+    result = async_algo.run(train_data, val_data)
+    assert inspect.iscoroutine(result)
+
+    # Await the result
+    await result
+
+    assert async_algo.run_executed  # type: ignore
+    assert async_algo.run_train == train_data  # type: ignore
+    assert async_algo.run_val == val_data  # type: ignore
+
+
+def test_algorithm_with_none_datasets():
+    """Test that algorithm works with None datasets."""
+
+    @algo
+    def nullable_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """Algorithm that accepts None."""
+        nullable_algo.called_with_none = train_dataset is None and val_dataset is None  # type: ignore
+
+    nullable_algo(train_dataset=None, val_dataset=None)
+    assert nullable_algo.called_with_none  # type: ignore
+
+    # Also test via run method
+    nullable_algo.called_with_none = False  # type: ignore
+    nullable_algo.run()
+    assert nullable_algo.called_with_none  # type: ignore
+
+
+def test_multiple_algorithm_instances():
+    """Test that multiple decorated functions work independently."""
+
+    @algo
+    def algo1(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """First algorithm."""
+        algo1.count = getattr(algo1, "count", 0) + 1  # type: ignore
+
+    @algo
+    def algo2(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """Second algorithm."""
+        algo2.count = getattr(algo2, "count", 0) + 1  # type: ignore
+
+    algo1.count = 0  # type: ignore
+    algo2.count = 0  # type: ignore
+
+    algo1(train_dataset=None, val_dataset=None)
+    algo1(train_dataset=None, val_dataset=None)
+    algo2(train_dataset=None, val_dataset=None)
+
+    assert algo1.count == 2  # type: ignore
+    assert algo2.count == 1  # type: ignore
+
+
+def test_algorithm_base_algorithm_methods():
+    """Test that Algorithm methods are available."""
+
+    @algo
+    def test_algo(*, train_dataset: Optional[Dataset[Any]], val_dataset: Optional[Dataset[Any]]) -> None:
+        """Test algorithm."""
+        pass
+
+    # Should have all Algorithm methods
+    assert hasattr(test_algo, "set_trainer")
+    assert hasattr(test_algo, "get_trainer")
+    assert hasattr(test_algo, "set_llm_proxy")
+    assert hasattr(test_algo, "get_llm_proxy")
+    assert hasattr(test_algo, "set_adapter")
+    assert hasattr(test_algo, "get_adapter")
+    assert hasattr(test_algo, "set_store")
+    assert hasattr(test_algo, "get_store")
+    assert hasattr(test_algo, "get_initial_resources")
+    assert hasattr(test_algo, "set_initial_resources")
+
+
+# New tests for parameter injection and error handling
+
+
+def test_algorithm_without_datasets():
+    """Test that algorithms can be defined without train_dataset/val_dataset parameters."""
+
+    @algo
+    def no_dataset_algo(*, store: LightningStore) -> None:
+        """Algorithm that doesn't use datasets."""
+        no_dataset_algo.store_passed = store  # type: ignore
+        no_dataset_algo.executed = True  # type: ignore
+
+    no_dataset_algo.executed = False  # type: ignore
+
+    # Set up the store
+    mock_store = MagicMock(spec=LightningStore)
+    no_dataset_algo.set_store(mock_store)
+
+    # Call run method without datasets
+    no_dataset_algo.run()
+
+    assert no_dataset_algo.executed  # type: ignore
+    assert no_dataset_algo.store_passed == mock_store  # type: ignore
+
+
+def test_algorithm_raises_error_on_unsupported_train_dataset():
+    """Test that TypeError is raised when train_dataset is provided but not supported."""
+
+    @algo
+    def no_train_algo(*, val_dataset: Optional[Dataset[Any]]) -> None:
+        """Algorithm that only accepts val_dataset."""
+        pass
+
+    # Providing train_dataset should raise TypeError
+    with pytest.raises(TypeError, match="train_dataset is provided but not supported"):
+        no_train_algo.run(train_dataset=["data"], val_dataset=None)
+
+
+def test_algorithm_raises_error_on_unsupported_val_dataset():
+    """Test that TypeError is raised when val_dataset is provided but not supported."""
+
+    @algo
+    def no_val_algo(*, train_dataset: Optional[Dataset[Any]]) -> None:
+        """Algorithm that only accepts train_dataset."""
+        pass
+
+    # Providing val_dataset should raise TypeError
+    with pytest.raises(TypeError, match="val_dataset is provided but not supported"):
+        no_val_algo.run(train_dataset=None, val_dataset=["data"])
+
+
+def test_algorithm_with_all_injected_parameters():
+    """Test that all injectable parameters (store, adapter, llm_proxy, initial_resources) work."""
+
+    @algo
+    def full_algo(
+        *,
+        store: LightningStore,
+        adapter: Any,
+        llm_proxy: Optional[Any] = None,
+        initial_resources: Optional[Any] = None,
+        train_dataset: Optional[Dataset[Any]],
+        val_dataset: Optional[Dataset[Any]],
+    ) -> None:
+        """Algorithm with all injectable parameters."""
+        full_algo.store = store  # type: ignore
+        full_algo.adapter = adapter  # type: ignore
+        full_algo.llm_proxy = llm_proxy  # type: ignore
+        full_algo.initial_resources = initial_resources  # type: ignore
+        full_algo.train = train_dataset  # type: ignore
+        full_algo.val = val_dataset  # type: ignore
+
+    # Set up all dependencies
+    mock_store = MagicMock(spec=LightningStore)
+    mock_adapter = MagicMock()
+    mock_llm_proxy = MagicMock()
+    mock_resources = MagicMock()
+
+    full_algo.set_store(mock_store)
+    full_algo.set_adapter(mock_adapter)
+    full_algo.set_llm_proxy(mock_llm_proxy)
+    full_algo.set_initial_resources(mock_resources)
+
+    train_data = ["train"]
+    val_data = ["val"]
+
+    # Run the algorithm
+    full_algo.run(train_data, val_data)
+
+    # Verify all parameters were injected correctly
+    assert full_algo.store == mock_store  # type: ignore
+    assert full_algo.adapter == mock_adapter  # type: ignore
+    assert full_algo.llm_proxy == mock_llm_proxy  # type: ignore
+    assert full_algo.initial_resources == mock_resources  # type: ignore
+    assert full_algo.train == train_data  # type: ignore
+    assert full_algo.val == val_data  # type: ignore
+
+
+def test_algorithm_with_only_store():
+    """Test algorithm that only uses the store parameter."""
+
+    @algo
+    def store_only_algo(*, store: LightningStore) -> None:
+        """Algorithm that only needs store."""
+        store_only_algo.got_store = True  # type: ignore
+        store_only_algo.store_value = store  # type: ignore
+
+    store_only_algo.got_store = False  # type: ignore
+
+    mock_store = MagicMock(spec=LightningStore)
+    store_only_algo.set_store(mock_store)
+
+    # Should work without any datasets
+    store_only_algo.run()
+
+    assert store_only_algo.got_store  # type: ignore
+    assert store_only_algo.store_value == mock_store  # type: ignore
+
+
+@pytest.mark.asyncio
+async def test_async_algorithm_with_injected_parameters():
+    """Test that async algorithms also support parameter injection."""
+
+    @algo
+    async def async_full_algo(
+        *,
+        store: LightningStore,
+        train_dataset: Optional[Dataset[Any]],
+    ) -> None:
+        """Async algorithm with injected parameters."""
+        async_full_algo.store = store  # type: ignore
+        async_full_algo.train = train_dataset  # type: ignore
+
+    mock_store = MagicMock(spec=LightningStore)
+    async_full_algo.set_store(mock_store)  # type: ignore
+
+    train_data = ["async-train"]
+    await async_full_algo.run(train_data)  # type: ignore
+
+    assert async_full_algo.store == mock_store  # type: ignore
+    assert async_full_algo.train == train_data  # type: ignore

tests/algorithm/test_utils.py

@@ -0,0 +1,115 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from __future__ import annotations
+
+from typing import Any, cast
+from unittest.mock import MagicMock
+
+import pytest
+
+from agentlightning.algorithm.base import Algorithm
+from agentlightning.algorithm.utils import with_llm_proxy, with_store
+from agentlightning.llm_proxy import LLMProxy
+from agentlightning.store.base import LightningStore
+
+
+class _BaseAlgorithm(Algorithm):
+    def run(self, *args: Any, **kwargs: Any) -> None:
+        """Satisfy the abstract interface without invoking training logic."""
+        return None
+
+
+class _StubLLMProxy:
+    """Test double that tracks lifecycle calls."""
+
+    def __init__(self) -> None:
+        self.start_calls = 0
+        self.stop_calls = 0
+        self.running = False
+
+    def is_running(self) -> bool:
+        return self.running
+
+    async def start(self) -> None:
+        self.start_calls += 1
+        self.running = True
+
+    async def stop(self) -> None:
+        self.stop_calls += 1
+        self.running = False
+
+
+@pytest.mark.asyncio
+async def test_with_store_injects_store_argument():
+    class StoreAlgorithm(_BaseAlgorithm):
+        @with_store
+        async def record_store(self, store: LightningStore, payload: str) -> None:
+            self.seen_store = store  # type: ignore[attr-defined]
+            self.seen_payload = payload  # type: ignore[attr-defined]
+
+    algorithm = StoreAlgorithm()
+    fake_store = MagicMock(spec=LightningStore)
+    algorithm.set_store(fake_store)
+
+    await algorithm.record_store("batch-1")
+
+    assert algorithm.seen_store is fake_store  # type: ignore[attr-defined]
+    assert algorithm.seen_payload == "batch-1"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_allows_optional_injection():
+    class OptionalProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def record_proxy(self, llm_proxy: LLMProxy | None, marker: str) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+            self.marker = marker  # type: ignore[attr-defined]
+
+    algorithm = OptionalProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    await algorithm.record_proxy("optional")
+
+    assert algorithm.seen_proxy is None  # type: ignore[attr-defined]
+    assert algorithm.marker == "optional"  # type: ignore[attr-defined]
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_required_raises_when_missing():
+    class RequiredProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy(required=True)
+        async def record_proxy(self, llm_proxy: LLMProxy) -> None:
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = RequiredProxyAlgorithm()
+    algorithm.set_llm_proxy(None)
+
+    with pytest.raises(ValueError):
+        await algorithm.record_proxy()
+
+
+@pytest.mark.asyncio
+async def test_with_llm_proxy_auto_start_and_stop():
+    class AutoProxyAlgorithm(_BaseAlgorithm):
+        @with_llm_proxy()
+        async def use_proxy(self, llm_proxy: LLMProxy | None) -> None:
+            if llm_proxy is None:
+                raise AssertionError("LLM proxy should be injected")
+            self.seen_proxy = llm_proxy  # type: ignore[attr-defined]
+
+    algorithm = AutoProxyAlgorithm()
+    proxy = _StubLLMProxy()
+    algorithm.set_llm_proxy(cast(LLMProxy, proxy))
+
+    await algorithm.use_proxy()
+
+    assert algorithm.seen_proxy is proxy  # type: ignore[attr-defined]
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1
+
+    # When already running, no extra start/stop should be requested.
+    proxy.running = True
+    await algorithm.use_proxy()
+
+    assert proxy.start_calls == 1
+    assert proxy.stop_calls == 1