fix(docker): allocate free host ports for sandboxes instead of deriving them from the flow id

Primary sandbox host ports were computed as
28000 + (flowID*containerPortsNumber + i) % limitContainerPortsNumber, so any two
flows whose ids are congruent mod (limitContainerPortsNumber/containerPortsNumber)
mapped to the same host ports. The second container then failed to start with
"port is already allocated" while its flow row lingered in "created" — and the
collision also crossed compose stacks sharing one docker daemon (e.g. flow 2 vs
flow 90002 both on 28004/28005).

Reserve free OS ports at bind time and hold the reservations open until docker
takes over, so two concurrent flows can never pick the same port. The agent
prompt now reads the container's actually-bound host ports back from docker
(correct after a restart too), keeping the deterministic set only as a
host-network / inspect-failure fallback. On a container-start failure the flow is
now marked failed instead of left stuck in "created".

Regression test: two flows N and N+period now get distinct free ports.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Sergey Kozyrenko
2026-07-20 15:39:51 +07:00
co-authored by Claude Fable 5
parent 2d1e677e75
commit eaf194cb89
9 changed files with 260 additions and 2 deletions
+17
View File
@@ -4,6 +4,7 @@ import (
"context"
"errors"
"fmt"
"time"
"pentagi/pkg/database"
"pentagi/pkg/graph/subscriptions"
@@ -49,6 +50,22 @@ type SubtaskContext struct {
TaskContext
}
// markFlowFailed records a flow as failed after its resources couldn't be
// prepared. It uses a detached context so the write still lands when the
// caller's context was the thing that got canceled.
func markFlowFailed(db database.Querier, flowID int64) {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if _, err := db.UpdateFlowStatus(ctx, database.UpdateFlowStatusParams{
Status: database.FlowStatusFailed,
ID: flowID,
}); err != nil {
logrus.WithError(err).WithField("flow_id", flowID).
Error("failed to mark flow as failed after resource preparation error")
}
}
func wrapErrorEndSpan(ctx context.Context, span langfuse.Span, msg string, err error) error {
logrus.WithContext(ctx).WithError(err).Error(msg)
err = fmt.Errorf("%s: %w", msg, err)
+3
View File
@@ -272,6 +272,9 @@ func NewFlowWorker(
}
if err := executor.Prepare(ctx); err != nil {
// The flow row is already persisted; a container that never started
// leaves it visibly failed instead of stuck in "created".
markFlowFailed(fwc.db, flow.ID)
return nil, wrapErrorEndSpan(ctx, flowSpan, "failed to prepare flow resources", err)
}
+98 -1
View File
@@ -7,10 +7,12 @@ import (
"fmt"
"hash/crc32"
"io"
"net"
"os"
"path"
"path/filepath"
"slices"
"strconv"
"strings"
"sync"
@@ -76,12 +78,22 @@ type DockerClient interface {
ContainerExecInspect(ctx context.Context, execID string) (container.ExecInspect, error)
ContainerStatPath(ctx context.Context, containerID string, path string) (container.PathStat, error)
ListContainerDir(ctx context.Context, containerID string, dirPath string) (ContainerDirListing, error)
// GetFlowContainerPorts returns the host ports actually bound to the flow's
// primary container (read from the live container, so it stays correct after
// a process restart and reflects the dynamically allocated ports).
GetFlowContainerPorts(ctx context.Context, flowID int64) ([]int, error)
CopyToContainer(ctx context.Context, containerID string, dstPath string, content io.Reader, options container.CopyToContainerOptions) error
CopyFromContainer(ctx context.Context, containerID string, srcPath string) (io.ReadCloser, container.PathStat, error)
Cleanup(ctx context.Context) error
GetDefaultImage() string
}
// GetPrimaryContainerPorts is the deterministic advisory port set for a flow.
// It is used only as a fallback for the prompt text (host-network mode, or when
// the live container can't be inspected) — the actual bridge-mode host bindings
// come from reserveFreePorts, since this derivation wraps modulo
// limitContainerPortsNumber and therefore collides for flow ids that are
// congruent mod (limitContainerPortsNumber/containerPortsNumber).
func GetPrimaryContainerPorts(flowID int64) []int {
ports := make([]int, containerPortsNumber)
for i := 0; i < containerPortsNumber; i++ {
@@ -91,6 +103,73 @@ func GetPrimaryContainerPorts(flowID int64) []int {
return ports
}
// reserveFreePorts asks the OS for n free TCP ports on bindIP and keeps each
// listener open, so two concurrent reservations can never hand back the same
// port. The caller closes the returned listeners immediately before Docker
// binds the ports (the only remaining race window is against a non-pentagi
// process, which the deterministic derivation could not avoid either).
func reserveFreePorts(bindIP string, n int) ([]int, []io.Closer, error) {
ports := make([]int, 0, n)
closers := make([]io.Closer, 0, n)
closeAll := func() {
for _, c := range closers {
_ = c.Close()
}
}
for i := 0; i < n; i++ {
listener, err := net.Listen("tcp", net.JoinHostPort(bindIP, "0"))
if err != nil {
closeAll()
return nil, nil, fmt.Errorf("failed to reserve free host port: %w", err)
}
closers = append(closers, listener)
ports = append(ports, listener.Addr().(*net.TCPAddr).Port)
}
return ports, closers, nil
}
func (dc *dockerClient) GetFlowContainerPorts(ctx context.Context, flowID int64) ([]int, error) {
cnt, err := dc.db.GetFlowPrimaryContainer(ctx, flowID)
if err != nil {
return nil, fmt.Errorf("failed to get primary container for flow %d: %w", flowID, err)
}
if !cnt.LocalID.Valid || cnt.LocalID.String == "" {
return nil, fmt.Errorf("primary container for flow %d has no local id", flowID)
}
inspection, err := dc.client.ContainerInspect(ctx, cnt.LocalID.String)
if err != nil {
return nil, fmt.Errorf("failed to inspect container %s: %w", cnt.LocalID.String, err)
}
return hostPortsFromBindings(inspection.NetworkSettings.Ports), nil
}
// hostPortsFromBindings collects the distinct host ports from a container's port
// map, sorted for a stable prompt rendering.
func hostPortsFromBindings(portMap nat.PortMap) []int {
seen := make(map[int]struct{})
for _, bindings := range portMap {
for _, binding := range bindings {
port, err := strconv.Atoi(binding.HostPort)
if err != nil {
continue
}
seen[port] = struct{}{}
}
}
ports := make([]int, 0, len(seen))
for port := range seen {
ports = append(ports, port)
}
slices.Sort(ports)
return ports
}
func NewDockerClient(ctx context.Context, db database.Querier, cfg *config.Config) (DockerClient, error) {
cli, err := client.NewClientWithOpts(client.FromEnv)
if err != nil {
@@ -292,6 +371,9 @@ func (dc *dockerClient) RunContainer(
// Configure network mode and port bindings
var networkingConfig *network.NetworkingConfig
// Held open until just before ContainerCreate so a concurrent flow cannot
// reserve the same host ports.
var portReservations []io.Closer
if dc.network == "host" {
// Host network mode: container uses host network stack directly
// No port bindings needed as container has direct access to host interfaces
@@ -305,7 +387,17 @@ func (dc *dockerClient) RunContainer(
if config.ExposedPorts == nil {
config.ExposedPorts = nat.PortSet{}
}
for _, port := range GetPrimaryContainerPorts(flowID) {
// Allocate free host ports instead of deriving them from flowID: the
// derivation wraps and collides for far-apart flow ids (and across
// stacks sharing this daemon), which fails the bind with "port is
// already allocated".
hostPorts, reservations, err := reserveFreePorts(dc.publicIP, containerPortsNumber)
if err != nil {
defer updateContainerInfo(database.ContainerStatusFailed, "")
return database.Container{}, fmt.Errorf("failed to allocate host ports for container '%s': %w", containerName, err)
}
portReservations = reservations
for _, port := range hostPorts {
natPort := nat.Port(fmt.Sprintf("%d/tcp", port))
hostConfig.PortBindings[natPort] = []nat.PortBinding{
{
@@ -325,6 +417,11 @@ func (dc *dockerClient) RunContainer(
}
}
// Hand the reserved host ports to Docker: release them so its bind succeeds.
for _, reservation := range portReservations {
_ = reservation.Close()
}
resp, err := dc.client.ContainerCreate(ctx, config, hostConfig, networkingConfig, nil, containerName)
if err != nil {
if config.Image == dc.defImage {
+95
View File
@@ -0,0 +1,95 @@
package docker
import (
"io"
"testing"
"github.com/docker/go-connections/nat"
)
// TestDeterministicPortsCollide documents why container host ports are no longer
// derived from the flow id: the derivation wraps modulo
// limitContainerPortsNumber, so two flows whose ids differ by
// limitContainerPortsNumber/containerPortsNumber map to the same host ports and
// the second container's bind fails with "port is already allocated".
func TestDeterministicPortsCollide(t *testing.T) {
period := int64(limitContainerPortsNumber / containerPortsNumber)
a := GetPrimaryContainerPorts(2)
b := GetPrimaryContainerPorts(2 + period)
if !slicesEqual(a, b) {
t.Fatalf("expected ids 2 and %d to collide, got %v vs %v", 2+period, a, b)
}
// 90002 ≡ 2 (mod period) → both derive to 28004/28005.
if got := GetPrimaryContainerPorts(90002); !slicesEqual(got, GetPrimaryContainerPorts(2)) {
t.Fatalf("expected flow 90002 to collide with flow 2, got %v", got)
}
}
// TestReserveFreePortsCoexist is the regression: two flows created concurrently
// (ids far apart, e.g. N and N+1000) must get distinct host ports. Holding both
// reservations open at once is exactly what the RunContainer path does until it
// hands the ports to Docker, so disjointness here proves the collision is gone.
func TestReserveFreePortsCoexist(t *testing.T) {
portsA, closersA, err := reserveFreePorts("127.0.0.1", containerPortsNumber)
if err != nil {
t.Fatalf("reserve A: %v", err)
}
defer closeAll(closersA)
portsB, closersB, err := reserveFreePorts("127.0.0.1", containerPortsNumber)
if err != nil {
t.Fatalf("reserve B: %v", err)
}
defer closeAll(closersB)
if len(portsA) != containerPortsNumber || len(portsB) != containerPortsNumber {
t.Fatalf("expected %d ports each, got %v and %v", containerPortsNumber, portsA, portsB)
}
all := append(append([]int{}, portsA...), portsB...)
seen := make(map[int]struct{}, len(all))
for _, p := range all {
if p <= 0 || p > 65535 {
t.Fatalf("port %d out of range", p)
}
if _, dup := seen[p]; dup {
t.Fatalf("two concurrent reservations collided on port %d (A=%v B=%v)", p, portsA, portsB)
}
seen[p] = struct{}{}
}
}
func TestHostPortsFromBindings(t *testing.T) {
portMap := nat.PortMap{
"28005/tcp": []nat.PortBinding{{HostIP: "127.0.0.1", HostPort: "40002"}},
"28004/tcp": []nat.PortBinding{{HostIP: "127.0.0.1", HostPort: "40001"}},
"9/tcp": nil, // no binding (host mode leaves these empty)
}
got := hostPortsFromBindings(portMap)
want := []int{40001, 40002}
if !slicesEqual(got, want) {
t.Fatalf("expected sorted host ports %v, got %v", want, got)
}
}
func closeAll(closers []io.Closer) {
for _, c := range closers {
_ = c.Close()
}
}
func slicesEqual(a, b []int) bool {
if len(a) != len(b) {
return false
}
for i := range a {
if a[i] != b[i] {
return false
}
}
return true
}
+33 -1
View File
@@ -860,8 +860,40 @@ func (fp *flowProvider) subtasksToMarkdown(subtasks []tools.SubtaskInfo) string
return buffer.String()
}
// containerPorts returns the primary container's actually bound host ports for
// the prompt (memoized, stable for the flow's lifetime). It falls back to the
// deterministic advisory set when the live ports can't be read — host-network
// mode publishes none, and a not-yet-inspectable container has none either.
func (fp *flowProvider) containerPorts() []int {
fp.mx.RLock()
cached := fp.cachedPorts
fp.mx.RUnlock()
if len(cached) > 0 {
return cached
}
fallback := docker.GetPrimaryContainerPorts(fp.flowID)
if fp.docker == nil {
return fallback
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
ports, err := fp.docker.GetFlowContainerPorts(ctx, fp.flowID)
if err != nil || len(ports) == 0 {
return fallback
}
fp.mx.Lock()
fp.cachedPorts = ports
fp.mx.Unlock()
return ports
}
func (fp *flowProvider) getContainerPortsDescription() string {
ports := docker.GetPrimaryContainerPorts(fp.flowID)
ports := fp.containerPorts()
var buffer strings.Builder
buffer.WriteString("**OOB Attack Infrastructure:**\n\n")
+4
View File
@@ -138,6 +138,10 @@ type flowProvider struct {
dataDir string
publicIP string
dockerNetwork string
docker docker.DockerClient
// cachedPorts memoizes the primary container's bound host ports (stable for
// the flow's lifetime) so the prompt doesn't inspect Docker on every render.
cachedPorts []int
callCounter *atomic.Int64
+4
View File
@@ -334,6 +334,7 @@ func (pc *providerController) NewFlowProvider(
dataDir: pc.cfg.DataDir,
publicIP: pc.publicIP,
dockerNetwork: pc.dockerNetwork,
docker: pc.docker,
callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)),
image: image,
title: title,
@@ -386,6 +387,7 @@ func (pc *providerController) LoadFlowProvider(
dataDir: pc.cfg.DataDir,
publicIP: pc.publicIP,
dockerNetwork: pc.dockerNetwork,
docker: pc.docker,
callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)),
image: image,
title: title,
@@ -483,6 +485,7 @@ func (pc *providerController) NewAssistantProvider(
dataDir: pc.cfg.DataDir,
publicIP: pc.publicIP,
dockerNetwork: pc.dockerNetwork,
docker: pc.docker,
callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)),
image: image,
title: title,
@@ -538,6 +541,7 @@ func (pc *providerController) LoadAssistantProvider(
dataDir: pc.cfg.DataDir,
publicIP: pc.publicIP,
dockerNetwork: pc.dockerNetwork,
docker: pc.docker,
callCounter: newAtomicInt64(pc.startCallNumber.Add(deltaCallCounter)),
image: image,
title: title,
@@ -1101,6 +1101,9 @@ func (f *fakeDockerClient) CopyFromContainer(_ context.Context, _ string, contai
}
return io.NopCloser(bytes.NewReader(f.copyFromBody)), f.copyFromStat, nil
}
func (f *fakeDockerClient) GetFlowContainerPorts(_ context.Context, _ int64) ([]int, error) {
return nil, nil
}
func (f *fakeDockerClient) Cleanup(_ context.Context) error { return nil }
func (f *fakeDockerClient) GetDefaultImage() string { return "test-image" }
+3
View File
@@ -99,6 +99,9 @@ func (m *contextAwareMockDockerClient) ContainerStatPath(_ context.Context, _ st
func (m *contextAwareMockDockerClient) ListContainerDir(_ context.Context, _ string, _ string) (docker.ContainerDirListing, error) {
return docker.ContainerDirListing{}, nil
}
func (m *contextAwareMockDockerClient) GetFlowContainerPorts(_ context.Context, _ int64) ([]int, error) {
return nil, nil
}
func (m *contextAwareMockDockerClient) ContainerExecInspect(_ context.Context, _ string) (container.ExecInspect, error) {
return m.inspectResp, nil
}