485 lines
12 KiB
Go
485 lines
12 KiB
Go
package agenttask
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strconv"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"iop/packages/go/agentguard"
|
|
)
|
|
|
|
type memoryStore struct {
|
|
mu sync.Mutex
|
|
revision uint64
|
|
state ManagerState
|
|
}
|
|
|
|
func newMemoryStore() *memoryStore {
|
|
return &memoryStore{state: ManagerState{SchemaVersion: currentSchemaVersion}}
|
|
}
|
|
|
|
func (s *memoryStore) Load(context.Context) (ManagerState, StateRevision, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return cloneState(s.state), StateRevision(strconv.FormatUint(s.revision, 10)), nil
|
|
}
|
|
|
|
func (s *memoryStore) CompareAndSwap(
|
|
_ context.Context,
|
|
expected StateRevision,
|
|
next ManagerState,
|
|
) (StateRevision, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if expected != StateRevision(strconv.FormatUint(s.revision, 10)) {
|
|
return "", ErrRevisionConflict
|
|
}
|
|
s.revision++
|
|
s.state = cloneState(next)
|
|
return StateRevision(strconv.FormatUint(s.revision, 10)), nil
|
|
}
|
|
|
|
func (s *memoryStore) edit(change func(*ManagerState)) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
next := cloneState(s.state)
|
|
change(&next)
|
|
s.revision++
|
|
s.state = next
|
|
}
|
|
|
|
func (s *memoryStore) snapshot() ManagerState {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
return cloneState(s.state)
|
|
}
|
|
|
|
type fixedClock struct {
|
|
now time.Time
|
|
}
|
|
|
|
func (c fixedClock) Now() time.Time { return c.now }
|
|
|
|
type fakeWorkflow struct {
|
|
mu sync.Mutex
|
|
snapshots map[ProjectID]ProjectWorkflowSnapshot
|
|
errors map[ProjectID]error
|
|
}
|
|
|
|
func (f *fakeWorkflow) RegisteredProjects(context.Context) ([]ProjectID, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
ids := make([]ProjectID, 0, len(f.snapshots)+len(f.errors))
|
|
seen := make(map[ProjectID]struct{})
|
|
for id := range f.snapshots {
|
|
ids = append(ids, id)
|
|
seen[id] = struct{}{}
|
|
}
|
|
for id := range f.errors {
|
|
if _, ok := seen[id]; !ok {
|
|
ids = append(ids, id)
|
|
}
|
|
}
|
|
return ids, nil
|
|
}
|
|
|
|
func (f *fakeWorkflow) Snapshot(
|
|
_ context.Context,
|
|
projectID ProjectID,
|
|
) (ProjectWorkflowSnapshot, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if err := f.errors[projectID]; err != nil {
|
|
return ProjectWorkflowSnapshot{}, err
|
|
}
|
|
snapshot, ok := f.snapshots[projectID]
|
|
if !ok {
|
|
return ProjectWorkflowSnapshot{}, fmt.Errorf("missing project %s", projectID)
|
|
}
|
|
return cloneWorkflow(snapshot), nil
|
|
}
|
|
|
|
type fakeSelector struct {
|
|
capacity int
|
|
}
|
|
|
|
func (s fakeSelector) Select(
|
|
_ context.Context,
|
|
request SelectionRequest,
|
|
) (ExecutionTarget, error) {
|
|
return ExecutionTarget{
|
|
ProviderID: "provider",
|
|
ModelID: "model",
|
|
ProfileID: "profile",
|
|
ProfileRevision: "profile-r1",
|
|
ConfigRevision: request.Project.Intent.ConfigRevision,
|
|
Capacity: s.capacity,
|
|
}, nil
|
|
}
|
|
|
|
type fakeIsolation struct {
|
|
t *testing.T
|
|
root string
|
|
mu sync.Mutex
|
|
baseRoots map[WorkspaceID]string
|
|
taskRoots map[string]string
|
|
preparations []string
|
|
}
|
|
|
|
func newFakeIsolation(t *testing.T) *fakeIsolation {
|
|
t.Helper()
|
|
return &fakeIsolation{
|
|
t: t, root: t.TempDir(),
|
|
baseRoots: make(map[WorkspaceID]string),
|
|
taskRoots: make(map[string]string),
|
|
}
|
|
}
|
|
|
|
func (f *fakeIsolation) Prepare(
|
|
_ context.Context,
|
|
request IsolationRequest,
|
|
) (PreparedIsolation, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
workspaceID := request.Project.WorkspaceID
|
|
baseRoot := f.baseRoots[workspaceID]
|
|
if baseRoot == "" {
|
|
baseRoot = filepath.Join(f.root, "base-"+string(workspaceID))
|
|
if err := os.MkdirAll(baseRoot, 0o755); err != nil {
|
|
f.t.Fatalf("create base root: %v", err)
|
|
}
|
|
f.baseRoots[workspaceID] = baseRoot
|
|
}
|
|
key := request.IdempotencyKey
|
|
taskRoot := f.taskRoots[key]
|
|
if taskRoot == "" {
|
|
taskRoot = filepath.Join(
|
|
f.root,
|
|
"task-"+string(request.Project.ProjectID)+"-"+string(request.Work.AttemptID),
|
|
)
|
|
if err := os.MkdirAll(taskRoot, 0o755); err != nil {
|
|
f.t.Fatalf("create task root: %v", err)
|
|
}
|
|
f.taskRoots[key] = taskRoot
|
|
}
|
|
f.preparations = append(f.preparations, key)
|
|
return PreparedIsolation{
|
|
Grant: &agentguard.WorkspaceGrant{
|
|
ProjectID: string(request.Project.ProjectID), WorkspaceID: string(workspaceID),
|
|
Root: baseRoot, Revision: string(request.Project.Intent.GrantRevision),
|
|
},
|
|
Descriptor: &agentguard.IsolationDescriptor{
|
|
ID: request.IdempotencyKey, Revision: "isolation-r1",
|
|
Mode: request.Work.Unit.IsolationMode, BaseRoot: baseRoot, TaskRoot: taskRoot,
|
|
WorkingDir: taskRoot, WritableRoots: []string{taskRoot},
|
|
PinnedBaseRevision: "base-r1", EnforcesWritableRoots: true,
|
|
},
|
|
Profile: agentguard.ProviderProfile{
|
|
ProviderID: request.Target.ProviderID, ModelID: request.Target.ModelID,
|
|
ProfileID: request.Target.ProfileID, Revision: request.Target.ProfileRevision,
|
|
Unattended: true, ApprovalBypass: true, WritableRootConfinement: true,
|
|
},
|
|
}, nil
|
|
}
|
|
|
|
type fakeInvoker struct {
|
|
mu sync.Mutex
|
|
results map[string]Submission
|
|
actualCalls []DispatchRequest
|
|
active int
|
|
maxActive int
|
|
delays map[WorkUnitID]time.Duration
|
|
}
|
|
|
|
func newFakeInvoker() *fakeInvoker {
|
|
return &fakeInvoker{
|
|
results: make(map[string]Submission),
|
|
delays: make(map[WorkUnitID]time.Duration),
|
|
}
|
|
}
|
|
|
|
func (f *fakeInvoker) Invoke(
|
|
ctx context.Context,
|
|
request DispatchRequest,
|
|
) (Submission, error) {
|
|
f.mu.Lock()
|
|
if previous, ok := f.results[request.IdempotencyKey]; ok {
|
|
f.mu.Unlock()
|
|
return previous, nil
|
|
}
|
|
if request.Permit == nil || request.Workspace.TaskRoot == request.Workspace.BaseRoot {
|
|
f.mu.Unlock()
|
|
return Submission{}, fmt.Errorf("unsafe dispatch request")
|
|
}
|
|
f.active++
|
|
if f.active > f.maxActive {
|
|
f.maxActive = f.active
|
|
}
|
|
delay := f.delays[request.Work.Unit.ID]
|
|
f.mu.Unlock()
|
|
if delay > 0 {
|
|
select {
|
|
case <-ctx.Done():
|
|
f.mu.Lock()
|
|
f.active--
|
|
f.mu.Unlock()
|
|
return Submission{}, ctx.Err()
|
|
case <-time.After(delay):
|
|
}
|
|
}
|
|
result := Submission{
|
|
ProjectID: request.Project.ProjectID, WorkUnitID: request.Work.Unit.ID,
|
|
AttemptID: request.Work.AttemptID,
|
|
ArtifactID: ArtifactID("artifact-" + string(request.Work.AttemptID)),
|
|
Ready: true,
|
|
}
|
|
f.mu.Lock()
|
|
f.active--
|
|
f.results[request.IdempotencyKey] = result
|
|
f.actualCalls = append(f.actualCalls, request)
|
|
f.mu.Unlock()
|
|
return result, nil
|
|
}
|
|
|
|
func (f *fakeInvoker) callCount() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return len(f.actualCalls)
|
|
}
|
|
|
|
func (f *fakeInvoker) maxConcurrency() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return f.maxActive
|
|
}
|
|
|
|
func (f *fakeInvoker) activeCount() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return f.active
|
|
}
|
|
|
|
func (f *fakeInvoker) roots() []string {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
roots := make([]string, 0, len(f.actualCalls))
|
|
for _, request := range f.actualCalls {
|
|
roots = append(roots, request.Workspace.TaskRoot)
|
|
}
|
|
return roots
|
|
}
|
|
|
|
type fakeReviewer struct {
|
|
mu sync.Mutex
|
|
sequences map[WorkUnitID][]ReviewVerdict
|
|
results map[string]ReviewResult
|
|
actualCalls []ReviewRequest
|
|
}
|
|
|
|
func newFakeReviewer() *fakeReviewer {
|
|
return &fakeReviewer{
|
|
sequences: make(map[WorkUnitID][]ReviewVerdict),
|
|
results: make(map[string]ReviewResult),
|
|
}
|
|
}
|
|
|
|
func (f *fakeReviewer) Review(
|
|
_ context.Context,
|
|
request ReviewRequest,
|
|
) (ReviewResult, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if previous, ok := f.results[request.IdempotencyKey]; ok {
|
|
return previous, nil
|
|
}
|
|
sequence := f.sequences[request.Work.Unit.ID]
|
|
index := int(request.Work.Attempt) - 1
|
|
verdict := ReviewVerdictPass
|
|
if index >= 0 && index < len(sequence) {
|
|
verdict = sequence[index]
|
|
}
|
|
result := ReviewResult{
|
|
ProjectID: request.Project.ProjectID, WorkUnitID: request.Work.Unit.ID,
|
|
AttemptID: request.Work.AttemptID, ArtifactID: request.Submission.ArtifactID,
|
|
Verdict: verdict, Message: string(verdict),
|
|
}
|
|
switch verdict {
|
|
case ReviewVerdictPass:
|
|
result.ChangeSet = &ChangeSetIdentity{
|
|
ID: ChangeSetID("change-" + string(request.Work.AttemptID)),
|
|
Revision: "change-r1", ArtifactID: request.Submission.ArtifactID,
|
|
}
|
|
case ReviewVerdictWarn, ReviewVerdictFail:
|
|
result.Rework = true
|
|
}
|
|
f.results[request.IdempotencyKey] = result
|
|
f.actualCalls = append(f.actualCalls, request)
|
|
return result, nil
|
|
}
|
|
|
|
func (f *fakeReviewer) callCount() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return len(f.actualCalls)
|
|
}
|
|
|
|
type fakeIntegrator struct {
|
|
mu sync.Mutex
|
|
outcomes map[WorkUnitID]IntegrationOutcome
|
|
results map[string]IntegrationResult
|
|
actualCalls []IntegrationRequest
|
|
}
|
|
|
|
func newFakeIntegrator() *fakeIntegrator {
|
|
return &fakeIntegrator{
|
|
outcomes: make(map[WorkUnitID]IntegrationOutcome),
|
|
results: make(map[string]IntegrationResult),
|
|
}
|
|
}
|
|
|
|
func (f *fakeIntegrator) Integrate(
|
|
_ context.Context,
|
|
request IntegrationRequest,
|
|
) (IntegrationResult, error) {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
if previous, ok := f.results[request.IdempotencyKey]; ok {
|
|
return previous, nil
|
|
}
|
|
outcome := f.outcomes[request.Work.Unit.ID]
|
|
if outcome == "" {
|
|
outcome = IntegrationOutcomeIntegrated
|
|
}
|
|
result := IntegrationResult{
|
|
ProjectID: request.Project.ProjectID, WorkUnitID: request.Work.Unit.ID,
|
|
ChangeSet: request.ChangeSet, Ordinal: request.Ordinal, Attempt: request.Attempt,
|
|
Outcome: outcome, BeforeRevision: "before",
|
|
}
|
|
if outcome == IntegrationOutcomeIntegrated {
|
|
result.AfterRevision = "after"
|
|
} else {
|
|
result.Retained = true
|
|
result.Blocker = &Blocker{Code: BlockerIntegrationFailed, Message: "fixture blocker"}
|
|
}
|
|
f.results[request.IdempotencyKey] = result
|
|
f.actualCalls = append(f.actualCalls, request)
|
|
return result, nil
|
|
}
|
|
|
|
func (f *fakeIntegrator) ordinals() []DispatchOrdinal {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
out := make([]DispatchOrdinal, 0, len(f.actualCalls))
|
|
for _, request := range f.actualCalls {
|
|
out = append(out, request.Ordinal)
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (f *fakeIntegrator) callCount() int {
|
|
f.mu.Lock()
|
|
defer f.mu.Unlock()
|
|
return len(f.actualCalls)
|
|
}
|
|
|
|
type recordingEvents struct {
|
|
mu sync.Mutex
|
|
events []Event
|
|
}
|
|
|
|
func (r *recordingEvents) Emit(_ context.Context, event Event) error {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
r.events = append(r.events, event)
|
|
return nil
|
|
}
|
|
|
|
func (r *recordingEvents) snapshot() []Event {
|
|
r.mu.Lock()
|
|
defer r.mu.Unlock()
|
|
return append([]Event(nil), r.events...)
|
|
}
|
|
|
|
type managerHarness struct {
|
|
t *testing.T
|
|
store *memoryStore
|
|
workflow *fakeWorkflow
|
|
isolation *fakeIsolation
|
|
invoker *fakeInvoker
|
|
reviewer *fakeReviewer
|
|
integrator *fakeIntegrator
|
|
events *recordingEvents
|
|
manager *Manager
|
|
}
|
|
|
|
func newHarness(
|
|
t *testing.T,
|
|
snapshots map[ProjectID]ProjectWorkflowSnapshot,
|
|
capacity int,
|
|
) *managerHarness {
|
|
t.Helper()
|
|
store := newMemoryStore()
|
|
workflow := &fakeWorkflow{snapshots: snapshots, errors: make(map[ProjectID]error)}
|
|
isolation := newFakeIsolation(t)
|
|
invoker := newFakeInvoker()
|
|
reviewer := newFakeReviewer()
|
|
integrator := newFakeIntegrator()
|
|
events := &recordingEvents{}
|
|
manager, err := NewManager(
|
|
ManagerConfig{
|
|
OwnerID: "manager-1", LeaseDuration: time.Minute,
|
|
MaxReworkAttempts: 3,
|
|
},
|
|
fixedClock{now: time.Date(2026, 7, 28, 0, 0, 0, 0, time.UTC)},
|
|
store, workflow, fakeSelector{capacity: capacity}, isolation,
|
|
invoker, reviewer, integrator, events,
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("NewManager: %v", err)
|
|
}
|
|
return &managerHarness{
|
|
t: t, store: store, workflow: workflow, isolation: isolation,
|
|
invoker: invoker, reviewer: reviewer, integrator: integrator,
|
|
events: events, manager: manager,
|
|
}
|
|
}
|
|
|
|
func testSnapshot(
|
|
projectID ProjectID,
|
|
workspaceID WorkspaceID,
|
|
units ...WorkUnit,
|
|
) ProjectWorkflowSnapshot {
|
|
return ProjectWorkflowSnapshot{
|
|
ProjectID: projectID, WorkspaceID: workspaceID,
|
|
Revision: "workflow-r1", Units: units,
|
|
}
|
|
}
|
|
|
|
func testUnit(id WorkUnitID, kind WriteSetKind) WorkUnit {
|
|
return WorkUnit{
|
|
ID: id, MilestoneID: "milestone", WriteSetKind: kind,
|
|
IsolationMode: agentguard.IsolationModeOverlay,
|
|
}
|
|
}
|
|
|
|
func (h *managerHarness) start(
|
|
projectID ProjectID,
|
|
workspaceID WorkspaceID,
|
|
autoResume *bool,
|
|
) {
|
|
h.t.Helper()
|
|
err := h.manager.StartProject(context.Background(), StartRequest{
|
|
CommandID: CommandID("start-" + string(projectID)),
|
|
ProjectID: projectID, WorkspaceID: workspaceID, MilestoneID: "milestone",
|
|
WorkflowRevision: "workflow-r1", ConfigRevision: "config-r1",
|
|
GrantRevision: "grant-r1", AutoResumeInterrupted: autoResume,
|
|
})
|
|
if err != nil {
|
|
h.t.Fatalf("StartProject(%s): %v", projectID, err)
|
|
}
|
|
}
|