mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: DevUI: Add OpenAI Responses API proxy support + HIL for Workflows (#1737)
* DevUI: Add OpenAI Responses API proxy support with enhanced UI features This commit adds support for proxying requests to OpenAI's Responses API, allowing DevUI to route conversations to OpenAI models when configured to enable testing. Backend changes: - Add OpenAI proxy executor with conversation routing logic - Enhance event mapper to support OpenAI Responses API format - Extend server endpoints to handle OpenAI proxy mode - Update models with OpenAI-specific response types - Remove emojis from logging and CLI output for cleaner text Frontend changes: - Add settings modal with OpenAI proxy configuration UI - Enhance agent and workflow views with improved state management - Add new UI components (separator, switch) for settings - Update debug panel with better event filtering - Improve message renderers for OpenAI content types - Update types and API client for OpenAI integration * update ui, settings modal and workflow input form, add register cleanup hooks. * add workflow HIL support, user mode, other fixes * feat(devui): add human-in-the-loop (HIL) support with dynamic response schemas Implement HIL workflow support allowing workflows to pause for user input with dynamically generated JSON schemas based on response handler type hints. Key Features: - Automatic response schema extraction from @response_handler decorators - Dynamic form generation in UI based on Pydantic/dataclass response types - Checkpoint-based conversation storage for HIL requests/responses - Resume workflow execution after user provides HIL response Backend Changes: - Add extract_response_type_from_executor() to introspect response handlers - Enrich RequestInfoEvent with response_schema via _enrich_request_info_event_with_response_schema() - Map RequestInfoEvent to response.input.requested OpenAI event format - Store HIL responses in conversation history and restore checkpoints Frontend Changes: - Add HILInputModal component with SchemaFormRenderer for dynamic forms - Support Pydantic BaseModel and dataclass response types - Render enum fields as dropdowns, strings as text/textarea, numbers, booleans, arrays, objects - Display original request context alongside response form Testing: - Add tests for checkpoint storage (test_checkpoints.py) - Add schema generation tests for all input types (test_schema_generation.py) - Validate end-to-end HIL flow with spam workflow sample This enables workflows to seamlessly pause execution and request structured user input with type-safe, validated forms generated automatically from response type annotations. * improve HIL support, improve workflow execution view * ui updates * ui updates * improve HIL for workflows, add auth and view modes * update workflow * security improvements , ui fixes * fix mypy error * update loading spinner in ui --------- Co-authored-by: Mark Wallace <127216156+markwallace-microsoft@users.noreply.github.com>
This commit is contained in:
committed by
GitHub
Unverified
parent
85484c0259
commit
94eae24082
@@ -0,0 +1,443 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for checkpoint-as-conversation-items implementation."""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from agent_framework import (
|
||||
Executor,
|
||||
InMemoryCheckpointStorage,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
from agent_framework_devui._conversations import (
|
||||
CheckpointConversationManager,
|
||||
InMemoryConversationStore,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowTestData:
|
||||
"""Simple test data."""
|
||||
|
||||
value: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class WorkflowHILRequest:
|
||||
"""HIL request for testing."""
|
||||
|
||||
question: str
|
||||
|
||||
|
||||
class WorkflowTestExecutor(Executor):
|
||||
"""Test executor with HIL."""
|
||||
|
||||
@handler
|
||||
async def process(self, data: WorkflowTestData, ctx: WorkflowContext) -> None:
|
||||
"""Process data and request approval."""
|
||||
await ctx.set_executor_state({"data_value": data.value})
|
||||
|
||||
# Request HIL (checkpoint created here)
|
||||
await ctx.request_info(request_data=WorkflowHILRequest(question=f"Approve {data.value}?"), response_type=str)
|
||||
|
||||
@response_handler
|
||||
async def handle_response(
|
||||
self, original_request: WorkflowHILRequest, response: str, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle HIL response."""
|
||||
state = await ctx.get_executor_state() or {}
|
||||
value = state.get("data_value", "")
|
||||
await ctx.send_message(f"{value}_approved" if response.lower() == "yes" else f"{value}_rejected")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def conversation_store():
|
||||
"""Create in-memory conversation store."""
|
||||
return InMemoryConversationStore()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def checkpoint_manager(conversation_store):
|
||||
"""Create checkpoint manager."""
|
||||
return CheckpointConversationManager(conversation_store)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_workflow():
|
||||
"""Create test workflow with checkpointing."""
|
||||
executor = WorkflowTestExecutor(id="test_executor")
|
||||
checkpoint_storage = InMemoryCheckpointStorage()
|
||||
|
||||
return (
|
||||
WorkflowBuilder(name="Test Workflow", description="Test checkpoint behavior")
|
||||
.set_start_executor(executor)
|
||||
.with_checkpointing(checkpoint_storage)
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
class TestCheckpointConversationManager:
|
||||
"""Test CheckpointConversationManager functionality - CONVERSATION-SCOPED."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_scoped_checkpoint_save(self, checkpoint_manager, test_workflow):
|
||||
"""Test checkpoint save in a specific conversation."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"conv_{entity_id}_test123"
|
||||
|
||||
# Create conversation first
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Create test checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"test": "data"}
|
||||
)
|
||||
|
||||
# Get checkpoint storage for this conversation and save
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_id = await storage.save_checkpoint(checkpoint)
|
||||
|
||||
assert checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
# Verify checkpoint stored in THIS conversation only
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
assert len(checkpoints) == 1
|
||||
assert checkpoints[0].checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversation_isolation(self, checkpoint_manager, test_workflow):
|
||||
"""Test that conversations are isolated - checkpoints don't leak between conversations."""
|
||||
entity_id = "test_entity"
|
||||
conv_a = f"conv_{entity_id}_aaa"
|
||||
conv_b = f"conv_{entity_id}_bbb"
|
||||
|
||||
# Create two conversations
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conv_a
|
||||
)
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conv_b
|
||||
)
|
||||
|
||||
# Save checkpoint to conversation A
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint_a = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"conversation": "A"},
|
||||
)
|
||||
storage_a = checkpoint_manager.get_checkpoint_storage(conv_a)
|
||||
await storage_a.save_checkpoint(checkpoint_a)
|
||||
|
||||
# Verify conversation A has checkpoint
|
||||
checkpoints_a = await storage_a.list_checkpoints()
|
||||
assert len(checkpoints_a) == 1
|
||||
|
||||
# Verify conversation B has NO checkpoints (isolation)
|
||||
storage_b = checkpoint_manager.get_checkpoint_storage(conv_b)
|
||||
checkpoints_b = await storage_b.list_checkpoints()
|
||||
assert len(checkpoints_b) == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_checkpoints_in_session(self, checkpoint_manager, test_workflow):
|
||||
"""Test listing checkpoints within a session."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_test456"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Save multiple checkpoints
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_ids = []
|
||||
for i in range(3):
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"iteration": i},
|
||||
)
|
||||
saved_id = await storage.save_checkpoint(checkpoint)
|
||||
checkpoint_ids.append(saved_id)
|
||||
|
||||
# List checkpoints using the storage
|
||||
checkpoints_list = await storage.list_checkpoints()
|
||||
assert len(checkpoints_list) == 3
|
||||
|
||||
# Verify all checkpoint IDs are present
|
||||
loaded_ids = [cp.checkpoint_id for cp in checkpoints_list]
|
||||
for saved_id in checkpoint_ids:
|
||||
assert saved_id in loaded_ids
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoints_appear_as_conversation_items(self, checkpoint_manager, test_workflow):
|
||||
"""Test that checkpoints appear as conversation items through the standard API."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_items_test"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Save multiple checkpoints
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
checkpoint_ids = []
|
||||
for i in range(2):
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=f"checkpoint_{i}",
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"iteration": i},
|
||||
)
|
||||
saved_id = await storage.save_checkpoint(checkpoint)
|
||||
checkpoint_ids.append(saved_id)
|
||||
|
||||
# List conversation items - should include checkpoints
|
||||
items, has_more = await checkpoint_manager.conversation_store.list_items(conversation_id)
|
||||
|
||||
# Filter for checkpoint items
|
||||
checkpoint_items = [item for item in items if (isinstance(item, dict) and item.get("type") == "checkpoint")]
|
||||
|
||||
# Verify we have the correct number of checkpoint items
|
||||
assert len(checkpoint_items) == 2, f"Expected 2 checkpoint items, got {len(checkpoint_items)}"
|
||||
|
||||
# Verify checkpoint items have correct structure
|
||||
for item in checkpoint_items:
|
||||
assert item.get("type") == "checkpoint"
|
||||
assert item.get("checkpoint_id") in checkpoint_ids
|
||||
assert item.get("workflow_id") == test_workflow.id
|
||||
assert "timestamp" in item
|
||||
assert item.get("id").startswith("checkpoint_") # ID format: checkpoint_{checkpoint_id}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_checkpoint_from_session(self, checkpoint_manager, test_workflow):
|
||||
"""Test loading checkpoint from a specific session."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_test789"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Create and save a checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
original_checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"test_key": "test_value"},
|
||||
)
|
||||
|
||||
# Save to this session
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
await storage.save_checkpoint(original_checkpoint)
|
||||
|
||||
# Load checkpoint from this session
|
||||
loaded_checkpoint = await storage.load_checkpoint(original_checkpoint.checkpoint_id)
|
||||
|
||||
assert loaded_checkpoint is not None
|
||||
assert loaded_checkpoint.checkpoint_id == original_checkpoint.checkpoint_id
|
||||
assert loaded_checkpoint.workflow_id == original_checkpoint.workflow_id
|
||||
assert loaded_checkpoint.shared_state == {"test_key": "test_value"}
|
||||
|
||||
|
||||
class TestCheckpointStorage:
|
||||
"""Test InMemoryCheckpointStorage per conversation - SESSION-SCOPED."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_storage_protocol(self, checkpoint_manager, test_workflow):
|
||||
"""Test that adapter implements CheckpointStorage protocol."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_adapter_test"
|
||||
|
||||
# Create session
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Get storage adapter for this session
|
||||
storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
|
||||
# Create test checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"test": "data"}
|
||||
)
|
||||
|
||||
# Test save_checkpoint
|
||||
checkpoint_id = await storage.save_checkpoint(checkpoint)
|
||||
assert checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
# Test load_checkpoint
|
||||
loaded = await storage.load_checkpoint(checkpoint_id)
|
||||
assert loaded is not None
|
||||
assert loaded.checkpoint_id == checkpoint_id
|
||||
|
||||
# Test list_checkpoint_ids
|
||||
ids = await storage.list_checkpoint_ids(workflow_id=test_workflow.id)
|
||||
assert checkpoint_id in ids
|
||||
|
||||
# Test list_checkpoints
|
||||
checkpoints_list = await storage.list_checkpoints(workflow_id=test_workflow.id)
|
||||
assert len(checkpoints_list) >= 1
|
||||
assert any(cp.checkpoint_id == checkpoint_id for cp in checkpoints_list)
|
||||
|
||||
|
||||
class TestIntegration:
|
||||
"""Integration tests for checkpoint workflow execution."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_manual_checkpoint_save_via_injected_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test manual checkpoint save via build-time storage injection."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test1"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Get checkpoint storage for this session
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
|
||||
# Set build-time storage (equivalent to .with_checkpointing() at build time)
|
||||
# Note: In production, DevUI uses runtime injection via run_stream() parameter
|
||||
if hasattr(test_workflow, "_runner") and hasattr(test_workflow._runner, "context"):
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Create and save a checkpoint via injected storage
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()), workflow_id=test_workflow.id, messages={}, shared_state={"injected": True}
|
||||
)
|
||||
await checkpoint_storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Verify checkpoint is accessible via storage (in this session)
|
||||
storage_checkpoints = await checkpoint_storage.list_checkpoints()
|
||||
assert len(storage_checkpoints) > 0
|
||||
assert storage_checkpoints[0].checkpoint_id == checkpoint.checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_roundtrip_via_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test checkpoint save/load roundtrip via storage adapter."""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test2"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Set build-time storage for testing
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Create checkpoint
|
||||
import uuid
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id=str(uuid.uuid4()),
|
||||
workflow_id=test_workflow.id,
|
||||
messages={},
|
||||
shared_state={"ready_to_resume": True},
|
||||
)
|
||||
checkpoint_id = await checkpoint_storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Verify checkpoint can be loaded for resume
|
||||
loaded = await checkpoint_storage.load_checkpoint(checkpoint_id)
|
||||
assert loaded is not None
|
||||
assert loaded.checkpoint_id == checkpoint_id
|
||||
assert loaded.shared_state == {"ready_to_resume": True}
|
||||
|
||||
# Verify checkpoint is accessible via storage (for UI to list checkpoints)
|
||||
checkpoints = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints) > 0
|
||||
assert checkpoints[0].checkpoint_id == checkpoint_id
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_workflow_auto_saves_checkpoints_to_injected_storage(self, checkpoint_manager, test_workflow):
|
||||
"""Test that workflows automatically save checkpoints to our conversation-backed storage.
|
||||
|
||||
This is the critical end-to-end test that verifies the entire checkpoint flow:
|
||||
1. Storage is set as build-time storage (simulates .with_checkpointing())
|
||||
2. Workflow runs and pauses at HIL point (IDLE_WITH_PENDING_REQUESTS status)
|
||||
3. Framework automatically saves checkpoint to our storage
|
||||
4. Checkpoint is accessible via manager for UI to list/resume
|
||||
|
||||
Note: In production, DevUI passes checkpoint_storage to run_stream() as runtime parameter.
|
||||
This test uses build-time injection to verify framework's checkpoint auto-save behavior.
|
||||
"""
|
||||
entity_id = "test_entity"
|
||||
conversation_id = f"session_{entity_id}_integration_test3"
|
||||
|
||||
# Create session conversation
|
||||
checkpoint_manager.conversation_store.create_conversation(
|
||||
metadata={"entity_id": entity_id, "type": "workflow_session"}, conversation_id=conversation_id
|
||||
)
|
||||
|
||||
# Set build-time storage to test automatic checkpoint saves
|
||||
checkpoint_storage = checkpoint_manager.get_checkpoint_storage(conversation_id)
|
||||
test_workflow._runner.context._checkpoint_storage = checkpoint_storage
|
||||
|
||||
# Verify no checkpoints initially
|
||||
checkpoints_before = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints_before) == 0
|
||||
|
||||
# Run workflow until it reaches IDLE_WITH_PENDING_REQUESTS (after checkpoint is created)
|
||||
saw_request_event = False
|
||||
async for event in test_workflow.run_stream(WorkflowTestData(value="test")):
|
||||
if hasattr(event, "__class__"):
|
||||
if event.__class__.__name__ == "RequestInfoEvent":
|
||||
saw_request_event = True
|
||||
# Wait for IDLE_WITH_PENDING_REQUESTS status (comes after checkpoint creation)
|
||||
is_status_event = event.__class__.__name__ == "WorkflowStatusEvent"
|
||||
has_pending_status = hasattr(event, "status") and "IDLE_WITH_PENDING_REQUESTS" in str(event.status)
|
||||
if is_status_event and has_pending_status:
|
||||
break
|
||||
|
||||
assert saw_request_event, "Test workflow should have emitted RequestInfoEvent"
|
||||
|
||||
# Verify checkpoint was AUTOMATICALLY saved to our storage by the framework
|
||||
checkpoints_after = await checkpoint_storage.list_checkpoints()
|
||||
assert len(checkpoints_after) > 0, "Workflow should have auto-saved checkpoint at HIL pause"
|
||||
|
||||
# Verify checkpoint has correct workflow_id
|
||||
checkpoint = checkpoints_after[0]
|
||||
assert checkpoint.workflow_id == test_workflow.id
|
||||
@@ -0,0 +1,365 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Tests for cleanup hook registration and execution."""
|
||||
|
||||
import asyncio
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role, TextContent
|
||||
|
||||
from agent_framework_devui import register_cleanup
|
||||
from agent_framework_devui._discovery import EntityDiscovery
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def cleanup_registry():
|
||||
"""Clear the cleanup registry before each test."""
|
||||
import agent_framework_devui
|
||||
|
||||
agent_framework_devui._cleanup_registry.clear()
|
||||
yield
|
||||
agent_framework_devui._cleanup_registry.clear()
|
||||
|
||||
|
||||
class MockAgent:
|
||||
"""Mock agent for testing."""
|
||||
|
||||
def __init__(self, name: str = "TestAgent"):
|
||||
self.id = f"test-{name.lower()}"
|
||||
self.name = name
|
||||
self.description = "Test agent for cleanup hooks"
|
||||
self.cleanup_called = False
|
||||
self.async_cleanup_called = False
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
"""Mock streaming run method."""
|
||||
yield AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, content=[TextContent(text="Test response")])],
|
||||
inner_messages=[],
|
||||
)
|
||||
|
||||
|
||||
class MockCredential:
|
||||
"""Mock credential object for testing cleanup."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
"""Mock async close method."""
|
||||
self.closed = True
|
||||
|
||||
|
||||
class MockSyncResource:
|
||||
"""Mock synchronous resource for testing cleanup."""
|
||||
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
def close(self):
|
||||
"""Mock sync close method."""
|
||||
self.closed = True
|
||||
|
||||
|
||||
# Test 1: Register single cleanup hook
|
||||
async def test_register_cleanup_single_hook():
|
||||
"""Test registering a single cleanup hook for an entity."""
|
||||
agent = MockAgent("SingleHook")
|
||||
credential = MockCredential()
|
||||
|
||||
# Register cleanup
|
||||
register_cleanup(agent, credential.close)
|
||||
|
||||
# Verify credential not closed yet
|
||||
assert not credential.closed
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get cleanup hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 1
|
||||
|
||||
# Execute hook
|
||||
await hooks[0]()
|
||||
assert credential.closed
|
||||
|
||||
|
||||
# Test 2: Register multiple cleanup hooks
|
||||
async def test_register_cleanup_multiple_hooks():
|
||||
"""Test registering multiple cleanup hooks for a single entity."""
|
||||
agent = MockAgent("MultipleHooks")
|
||||
credential1 = MockCredential()
|
||||
credential2 = MockCredential()
|
||||
sync_resource = MockSyncResource()
|
||||
|
||||
# Register multiple hooks at once
|
||||
register_cleanup(agent, credential1.close, credential2.close, sync_resource.close)
|
||||
|
||||
# Verify nothing closed yet
|
||||
assert not credential1.closed
|
||||
assert not credential2.closed
|
||||
assert not sync_resource.closed
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get and execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 3
|
||||
|
||||
# Execute all hooks
|
||||
for hook in hooks:
|
||||
if asyncio.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
|
||||
assert credential1.closed
|
||||
assert credential2.closed
|
||||
assert sync_resource.closed
|
||||
|
||||
|
||||
# Test 3: Register cleanup hooks incrementally
|
||||
async def test_register_cleanup_incremental():
|
||||
"""Test registering cleanup hooks in multiple calls."""
|
||||
agent = MockAgent("IncrementalHooks")
|
||||
credential1 = MockCredential()
|
||||
credential2 = MockCredential()
|
||||
|
||||
# Register hooks incrementally
|
||||
register_cleanup(agent, credential1.close)
|
||||
register_cleanup(agent, credential2.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Should have both hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
# Execute all hooks
|
||||
for hook in hooks:
|
||||
await hook()
|
||||
|
||||
assert credential1.closed
|
||||
assert credential2.closed
|
||||
|
||||
|
||||
# Test 4: Test with no cleanup hooks
|
||||
async def test_no_cleanup_hooks():
|
||||
"""Test entity without any cleanup hooks registered."""
|
||||
agent = MockAgent("NoHooks")
|
||||
|
||||
# Don't register any cleanup hooks
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Should return empty list
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 0
|
||||
|
||||
|
||||
# Test 5: Test cleanup with async and sync hooks mixed
|
||||
async def test_mixed_async_sync_hooks():
|
||||
"""Test that both async and sync cleanup hooks work together."""
|
||||
agent = MockAgent("MixedHooks")
|
||||
async_resource = MockCredential()
|
||||
sync_resource = MockSyncResource()
|
||||
|
||||
# Register both types
|
||||
register_cleanup(agent, async_resource.close, sync_resource.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get and execute hooks with proper async/sync handling
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
import inspect
|
||||
|
||||
for hook in hooks:
|
||||
if inspect.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
|
||||
assert async_resource.closed
|
||||
assert sync_resource.closed
|
||||
|
||||
|
||||
# Test 6: Test error handling in cleanup hooks
|
||||
async def test_cleanup_hook_error_handling():
|
||||
"""Test that errors in cleanup hooks don't break execution."""
|
||||
agent = MockAgent("ErrorHooks")
|
||||
credential = MockCredential()
|
||||
|
||||
def failing_hook():
|
||||
raise RuntimeError("Intentional error for testing")
|
||||
|
||||
# Register failing hook and valid hook
|
||||
register_cleanup(agent, failing_hook, credential.close)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Get hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 2
|
||||
|
||||
# Execute hooks with error handling (like _server.py does)
|
||||
import inspect
|
||||
|
||||
for hook in hooks:
|
||||
try:
|
||||
if inspect.iscoroutinefunction(hook):
|
||||
await hook()
|
||||
else:
|
||||
hook()
|
||||
except Exception:
|
||||
pass # Ignore errors like the server does
|
||||
|
||||
# Second hook should still execute despite first one failing
|
||||
await credential.close()
|
||||
assert credential.closed
|
||||
|
||||
|
||||
# Test 7: Test ValueError when no hooks provided
|
||||
def test_register_cleanup_no_hooks_error():
|
||||
"""Test that register_cleanup raises ValueError when no hooks provided."""
|
||||
agent = MockAgent("NoHooksError")
|
||||
|
||||
with pytest.raises(ValueError, match="At least one cleanup hook required"):
|
||||
register_cleanup(agent)
|
||||
|
||||
|
||||
# Test 8: Test file-based discovery with cleanup hooks
|
||||
async def test_cleanup_with_file_based_discovery():
|
||||
"""Test that cleanup hooks work with file-based entity discovery."""
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
temp_path = Path(temp_dir)
|
||||
|
||||
# Create agent directory
|
||||
agent_dir = temp_path / "test_agent"
|
||||
agent_dir.mkdir()
|
||||
|
||||
# Write agent module with cleanup registration
|
||||
agent_file = agent_dir / "__init__.py"
|
||||
agent_file.write_text("""
|
||||
from agent_framework import AgentRunResponse, ChatMessage, Role, TextContent
|
||||
from agent_framework_devui import register_cleanup
|
||||
|
||||
class MockCredential:
|
||||
def __init__(self):
|
||||
self.closed = False
|
||||
|
||||
async def close(self):
|
||||
self.closed = True
|
||||
|
||||
# Create credential and agent
|
||||
credential = MockCredential()
|
||||
|
||||
class TestAgent:
|
||||
id = "test-agent"
|
||||
name = "Test Agent"
|
||||
description = "Test agent with cleanup"
|
||||
|
||||
async def run_stream(self, messages=None, *, thread=None, **kwargs):
|
||||
yield AgentRunResponse(
|
||||
messages=[ChatMessage(role=Role.ASSISTANT, content=[TextContent(text="Test")])],
|
||||
inner_messages=[],
|
||||
)
|
||||
|
||||
agent = TestAgent()
|
||||
|
||||
# Register cleanup at module level
|
||||
register_cleanup(agent, credential.close)
|
||||
""")
|
||||
|
||||
# Discover entities
|
||||
discovery = EntityDiscovery(str(temp_path))
|
||||
await discovery.discover_entities()
|
||||
|
||||
# Load the entity (triggers module import)
|
||||
await discovery.load_entity("test_agent")
|
||||
|
||||
# Verify cleanup hooks were registered
|
||||
hooks = discovery.get_cleanup_hooks("test_agent")
|
||||
assert len(hooks) == 1
|
||||
|
||||
|
||||
# Test 9: Test cleanup execution order
|
||||
async def test_cleanup_execution_order():
|
||||
"""Test that cleanup hooks execute in registration order."""
|
||||
agent = MockAgent("OrderTest")
|
||||
execution_order = []
|
||||
|
||||
def hook1():
|
||||
execution_order.append(1)
|
||||
|
||||
def hook2():
|
||||
execution_order.append(2)
|
||||
|
||||
def hook3():
|
||||
execution_order.append(3)
|
||||
|
||||
# Register in specific order
|
||||
register_cleanup(agent, hook1, hook2, hook3)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
for hook in hooks:
|
||||
hook()
|
||||
|
||||
# Verify execution order
|
||||
assert execution_order == [1, 2, 3]
|
||||
|
||||
|
||||
# Test 10: Test custom cleanup logic
|
||||
async def test_custom_cleanup_logic():
|
||||
"""Test registering custom cleanup function with complex logic."""
|
||||
agent = MockAgent("CustomCleanup")
|
||||
cleanup_executed = False
|
||||
resources_closed = []
|
||||
|
||||
async def custom_cleanup():
|
||||
nonlocal cleanup_executed
|
||||
cleanup_executed = True
|
||||
resources_closed.append("credential")
|
||||
resources_closed.append("session")
|
||||
resources_closed.append("cache")
|
||||
|
||||
register_cleanup(agent, custom_cleanup)
|
||||
|
||||
# Simulate discovery and registration
|
||||
discovery = EntityDiscovery()
|
||||
entity_info = await discovery.create_entity_info_from_object(agent, entity_type="agent", source="in_memory")
|
||||
discovery.register_entity(entity_info.id, entity_info, agent)
|
||||
|
||||
# Execute hooks
|
||||
hooks = discovery.get_cleanup_hooks(entity_info.id)
|
||||
assert len(hooks) == 1
|
||||
|
||||
await hooks[0]()
|
||||
|
||||
assert cleanup_executed
|
||||
assert resources_closed == ["credential", "session", "cache"]
|
||||
@@ -415,6 +415,56 @@ async def test_executor_action_events(mapper: MessageMapper, test_request: Agent
|
||||
assert "Executor failed" in str(events[0].item["error"]["message"])
|
||||
|
||||
|
||||
async def test_magentic_agent_delta_creates_message_container(
|
||||
mapper: MessageMapper, test_request: AgentFrameworkRequest
|
||||
) -> None:
|
||||
"""Test that MagenticAgentDeltaEvent creates message containers (Option A implementation)."""
|
||||
|
||||
# Create mock MagenticAgentDeltaEvent that mimics the real class
|
||||
from dataclasses import dataclass
|
||||
|
||||
try:
|
||||
from agent_framework import WorkflowEvent
|
||||
|
||||
@dataclass
|
||||
class MagenticAgentDeltaEvent(WorkflowEvent): # Inherit from WorkflowEvent
|
||||
agent_id: str
|
||||
text: str | None = None
|
||||
|
||||
except ImportError:
|
||||
# Fallback if WorkflowEvent is not available
|
||||
@dataclass
|
||||
class MagenticAgentDeltaEvent: # Use the expected name directly
|
||||
agent_id: str
|
||||
text: str | None = None
|
||||
|
||||
# First delta should create message container
|
||||
first_delta = MagenticAgentDeltaEvent(agent_id="test_agent", text="Hello ")
|
||||
events = await mapper.convert_event(first_delta, test_request)
|
||||
|
||||
# Should emit 3 events: message container, content part, and text delta
|
||||
assert len(events) == 3
|
||||
assert events[0].type == "response.output_item.added"
|
||||
assert events[0].item.type == "message" # Message, not executor_action!
|
||||
assert events[0].item.metadata["agent_id"] == "test_agent"
|
||||
assert events[0].item.metadata["source"] == "magentic"
|
||||
message_id = events[0].item.id
|
||||
|
||||
# Check text delta references the message ID
|
||||
assert events[2].type == "response.output_text.delta"
|
||||
assert events[2].item_id == message_id
|
||||
assert events[2].delta == "Hello "
|
||||
|
||||
# Second delta should NOT create new container
|
||||
second_delta = MagenticAgentDeltaEvent(agent_id="test_agent", text="world!")
|
||||
events = await mapper.convert_event(second_delta, test_request)
|
||||
|
||||
# Only text delta, no new container
|
||||
assert len(events) == 1
|
||||
assert events[0].type == "response.output_text.delta"
|
||||
assert events[0].item_id == message_id
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_all_tests() -> None:
|
||||
|
||||
@@ -4,13 +4,14 @@
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
import pytest
|
||||
|
||||
# Add parent package to path
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
|
||||
from agent_framework_devui._utils import generate_input_schema
|
||||
from agent_framework_devui._utils import extract_response_type_from_executor, generate_input_schema
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -132,6 +133,99 @@ def test_schema_generation_error_handling():
|
||||
pass
|
||||
|
||||
|
||||
def test_extract_response_type_from_executor():
|
||||
"""Test extraction of response type from @response_handler methods."""
|
||||
try:
|
||||
from agent_framework import Executor, WorkflowContext, handler, response_handler
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
# Define test request and response types
|
||||
@dataclass
|
||||
class TestApprovalRequest:
|
||||
"""Test request for approval."""
|
||||
|
||||
prompt: str
|
||||
context: str
|
||||
|
||||
class TestDecision(BaseModel):
|
||||
"""Test decision response."""
|
||||
|
||||
decision: Literal["approve", "reject"] = Field(description="User's decision")
|
||||
reason: str = Field(description="Reason for decision", default="")
|
||||
|
||||
# Create test executor with @response_handler
|
||||
class TestExecutor(Executor):
|
||||
"""Test executor with response handler."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Regular handler to satisfy executor requirements."""
|
||||
# Request info that will be handled by response_handler
|
||||
request = TestApprovalRequest(prompt="Test", context="Test context")
|
||||
await ctx.request_info(request, TestDecision)
|
||||
|
||||
@response_handler
|
||||
async def handle_approval(
|
||||
self, original_request: TestApprovalRequest, response: TestDecision, ctx: WorkflowContext
|
||||
) -> None:
|
||||
"""Handle approval response."""
|
||||
pass
|
||||
|
||||
# Test extraction
|
||||
executor = TestExecutor()
|
||||
extracted_type = extract_response_type_from_executor(executor, TestApprovalRequest)
|
||||
|
||||
# Verify correct type was extracted
|
||||
assert extracted_type is not None, "Should extract response type from @response_handler"
|
||||
assert extracted_type == TestDecision, f"Expected TestDecision, got {extracted_type}"
|
||||
|
||||
# Test full schema generation pipeline
|
||||
schema = generate_input_schema(extracted_type)
|
||||
assert schema is not None
|
||||
assert isinstance(schema, dict)
|
||||
assert "properties" in schema
|
||||
assert "decision" in schema["properties"]
|
||||
assert "enum" in schema["properties"]["decision"]
|
||||
assert schema["properties"]["decision"]["enum"] == ["approve", "reject"]
|
||||
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Required dependencies not available: {e}")
|
||||
|
||||
|
||||
def test_extract_response_type_no_match():
|
||||
"""Test that extraction returns None when no matching handler exists."""
|
||||
try:
|
||||
from agent_framework import Executor, WorkflowContext, handler
|
||||
|
||||
@dataclass
|
||||
class UnmatchedRequest:
|
||||
"""Request type with no handler."""
|
||||
|
||||
data: str
|
||||
|
||||
class MinimalExecutor(Executor):
|
||||
"""Executor with a handler but no matching response_handler."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="minimal_executor")
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Regular handler."""
|
||||
pass
|
||||
|
||||
executor = MinimalExecutor()
|
||||
extracted_type = extract_response_type_from_executor(executor, UnmatchedRequest)
|
||||
|
||||
assert extracted_type is None, "Should return None when no matching handler exists"
|
||||
|
||||
except ImportError as e:
|
||||
pytest.skip(f"Required dependencies not available: {e}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner for manual execution
|
||||
pytest.main([__file__, "-v"])
|
||||
|
||||
@@ -241,6 +241,100 @@ async def test_multiple_credential_attributes() -> None:
|
||||
assert mock_cred2.close.called, "Async credential should be closed"
|
||||
|
||||
|
||||
def test_ui_mode_configuration():
|
||||
"""Test UI mode configuration."""
|
||||
dev_server = DevServer(mode="developer")
|
||||
assert dev_server.mode == "developer"
|
||||
|
||||
user_server = DevServer(mode="user")
|
||||
assert user_server.mode == "user"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_api_restrictions_in_user_mode():
|
||||
"""Test that developer APIs are restricted in user mode."""
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
# Create servers with different modes
|
||||
dev_server = DevServer(mode="developer")
|
||||
user_server = DevServer(mode="user")
|
||||
|
||||
dev_app = dev_server.create_app()
|
||||
user_app = user_server.create_app()
|
||||
|
||||
dev_client = TestClient(dev_app)
|
||||
user_client = TestClient(user_app)
|
||||
|
||||
# Test 1: Health endpoint should work in both modes
|
||||
assert dev_client.get("/health").status_code == 200
|
||||
assert user_client.get("/health").status_code == 200
|
||||
|
||||
# Test 2: Meta endpoint should reflect correct mode
|
||||
dev_meta = dev_client.get("/meta").json()
|
||||
assert dev_meta["ui_mode"] == "developer"
|
||||
|
||||
user_meta = user_client.get("/meta").json()
|
||||
assert user_meta["ui_mode"] == "user"
|
||||
|
||||
# Test 3: Entity listing should work in both modes
|
||||
assert dev_client.get("/v1/entities").status_code == 200
|
||||
assert user_client.get("/v1/entities").status_code == 200
|
||||
|
||||
# Test 4: Entity info should be restricted in user mode
|
||||
dev_response = dev_client.get("/v1/entities/test_agent/info")
|
||||
assert dev_response.status_code in [200, 404, 500] # Not 403
|
||||
|
||||
user_response = user_client.get("/v1/entities/test_agent/info")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
# FastAPI wraps HTTPException detail in 'detail' field
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert error is not None
|
||||
assert "developer mode" in error["message"].lower()
|
||||
assert error["code"] == "developer_mode_required"
|
||||
|
||||
# Test 5: Hot reload should be restricted in user mode
|
||||
dev_response = dev_client.post("/v1/entities/test_agent/reload")
|
||||
assert dev_response.status_code in [200, 404, 500] # Not 403
|
||||
|
||||
user_response = user_client.post("/v1/entities/test_agent/reload")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert "developer mode" in error["message"].lower()
|
||||
|
||||
# Test 6: Deployment endpoints should be restricted in user mode
|
||||
# List deployments (simplest test - no payload needed)
|
||||
user_response = user_client.get("/v1/deployments")
|
||||
assert user_response.status_code == 403
|
||||
error_data = user_response.json()
|
||||
error = error_data.get("detail", {}).get("error") or error_data.get("error")
|
||||
assert "developer mode" in error["message"].lower()
|
||||
|
||||
# Get deployment
|
||||
user_response = user_client.get("/v1/deployments/test-id")
|
||||
assert user_response.status_code == 403
|
||||
|
||||
# Delete deployment
|
||||
user_response = user_client.delete("/v1/deployments/test-id")
|
||||
assert user_response.status_code == 403
|
||||
|
||||
# Test 7: Conversation endpoints should work in both modes
|
||||
dev_response = dev_client.post("/v1/conversations", json={})
|
||||
assert dev_response.status_code == 200
|
||||
|
||||
user_response = user_client.post("/v1/conversations", json={})
|
||||
assert user_response.status_code == 200
|
||||
|
||||
# Test 8: Chat endpoint should work in both modes
|
||||
chat_payload = {"model": "test_agent", "input": "Hello"}
|
||||
dev_response = dev_client.post("/v1/responses", json=chat_payload)
|
||||
assert dev_response.status_code in [200, 404] # 404 if agent doesn't exist
|
||||
|
||||
user_response = user_client.post("/v1/responses", json=chat_payload)
|
||||
assert user_response.status_code in [200, 404]
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# Simple test runner
|
||||
async def run_tests():
|
||||
@@ -273,3 +367,44 @@ class WeatherAgent:
|
||||
await executor.execute_sync(request)
|
||||
|
||||
asyncio.run(run_tests())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_checkpoint_api_endpoints(test_entities_dir):
|
||||
"""Test checkpoint list and delete API endpoints."""
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
|
||||
server = DevServer(entities_dir=test_entities_dir)
|
||||
executor = await server._ensure_executor()
|
||||
|
||||
# Create a conversation
|
||||
conversation = executor.conversation_store.create_conversation(metadata={"name": "Test Session"})
|
||||
conv_id = conversation.id
|
||||
|
||||
# Get checkpoint storage and add a checkpoint
|
||||
storage = executor.checkpoint_manager.get_checkpoint_storage(conv_id)
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id="test_checkpoint_1",
|
||||
workflow_id="test_workflow",
|
||||
shared_state={"key": "value"},
|
||||
iteration_count=1,
|
||||
)
|
||||
await storage.save_checkpoint(checkpoint)
|
||||
|
||||
# Test list checkpoints endpoint
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
assert len(checkpoints) == 1
|
||||
assert checkpoints[0].checkpoint_id == "test_checkpoint_1"
|
||||
assert checkpoints[0].workflow_id == "test_workflow"
|
||||
|
||||
# Test delete checkpoint endpoint
|
||||
deleted = await storage.delete_checkpoint("test_checkpoint_1")
|
||||
assert deleted is True
|
||||
|
||||
# Verify checkpoint was deleted
|
||||
remaining = await storage.list_checkpoints()
|
||||
assert len(remaining) == 0
|
||||
|
||||
# Test delete non-existent checkpoint
|
||||
deleted = await storage.delete_checkpoint("nonexistent")
|
||||
assert deleted is False
|
||||
|
||||
Reference in New Issue
Block a user