Help us shape the future of Go! Take the 2026 Go Developer Survey

agentrpc

package
v0.5.8 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Nov 1, 2026 License: Apache-2.0 Imports: 24 Imported by: 0

Documentation

Overview

Package agentrpc implements the control-plane side of the agent gRPC protocol: it authenticates each in-pod agent by its per-task-instance token, serves the task specification, and records the state transitions the agent reports.

Index

Constants

View Source 
const (
	// ScopingPermissive delivers the whole tenant vault when a DAG declares
	// nothing (today's behavior) and warns — but still delivers the whole vault —
	// when a DAG declares a narrower set. Subsetting is reserved for enforce, so no
	// already-declaring pipeline loses access. It is the default for this shipment.
	ScopingPermissive = "permissive"
	// ScopingEnforce delivers ONLY the declared subset (empty declaration →
	// nothing), resolved server-side and filtered in the query.
	ScopingEnforce = "enforce"
	// ScopingOff disables scoping entirely: the whole tenant vault, no warn.
	ScopingOff = "off"
)

Secret scoping policy (ADR 0055 D9), operator-set, NEVER author-settable.

View Source 
const (
	// LivenessObserve logs a structured warn + records a would-have-denied audit
	// event when the caller's TI is not live, but STILL delivers. It is the
	// default: no behavior change, the observe half of the warn→enforce arc.
	LivenessObserve = "observe"
	// LivenessEnforce denies with codes.PermissionDenied when the caller's TI is
	// not live. It is the operator's later flip, after an observe period.
	LivenessEnforce = "enforce"
)

Secret-path liveness gate modes (ADR 0055 E2). The gate consults the read-only task-instance liveness predicate before serving secrets so a token whose task instance is no longer live stops resolving them.

View Source 
const DecisionWarmWorkerNoSourceDelivery = "warm_worker_no_source_delivery"

DecisionWarmWorkerNoSourceDelivery is the decision type metered, once per worker, when a source-mode attempt passes over a free warm worker that did not register with source_delivery (ADR 0095): its agent predates it, and the attempt takes a dedicated pod unless another worker can run it.

View Source 
const DecisionWarmWorkerRefusedNoSweep = "warm_worker_refused_no_sweep"

DecisionWarmWorkerRefusedNoSweep is the decision type metered when a warm worker is refused because it does not confirm the attempt home sweep (SweepRefuse).

View Source 
const DecisionWarmWorkerSweepUnconfirmed = "warm_worker_sweep_unconfirmed"

DecisionWarmWorkerSweepUnconfirmed is the decision type metered when a warm worker that does not confirm the attempt home sweep is accepted (SweepWarn).

View Source 
const ReasonAttemptHomeSweepUnconfirmed = "attempt_home_sweep_unconfirmed"

ReasonAttemptHomeSweepUnconfirmed is the stable reason logged for a warm worker that does not confirm the sweep, and carried by its refusal.

Variables

View Source 
var ErrStaleReport = errors.New("task state report did not apply: the task instance already moved on")

ErrStaleReport is returned by Store.ReportState when the report did not apply because the task instance had already moved on — a reaper settled it, or a retry advanced past the attempt the reporting agent was dispatched for. It is not a failure of the RPC: the agent did nothing wrong and must not retry, so the handler acknowledges and logs. Declared here rather than in the storage package because storage implements this interface, not the reverse.

View Source 
var ErrWorkerAlreadyRegistered = errors.New("warm worker identity already has a live registration")

ErrWorkerAlreadyRegistered is returned by Register when the identity already has a live registration: its stream is still connected and has heartbeated within workerLivenessGrace.

Functions

func KeepaliveEnforcementPolicy added in v0.5.4 

func KeepaliveEnforcementPolicy() keepalive.EnforcementPolicy

KeepaliveEnforcementPolicy is the client-ping policy the agent gRPC server enforces. It permits a ping every 10 seconds, with or without an active stream, so neither a warm worker's idle assignment stream nor an agent's client keepalive (#946) is closed as abusive ("too_many_pings"). It only relaxes the gRPC default (5 minutes, streams only).

The agent's own keepalive interval must stay at or above MinTime; the agent's tests pin that.

func RecoveryStreamInterceptor 

func RecoveryStreamInterceptor(logger *slog.Logger) grpc.StreamServerInterceptor

RecoveryStreamInterceptor is RecoveryUnaryInterceptor for streaming handlers (e.g. log streaming), so a panic mid-stream returns Internal instead of crashing the control plane.

func RecoveryUnaryInterceptor 

func RecoveryUnaryInterceptor(logger *slog.Logger) grpc.UnaryServerInterceptor

RecoveryUnaryInterceptor recovers panics in unary RPC handlers so a single malformed or unexpected request from an agent cannot crash the control plane. The panic is logged with its stack and translated to a gRPC Internal error; the server keeps serving. It should be the outermost interceptor so it also covers any later interceptor.

Types

type AgentTokenMinter 

type AgentTokenMinter interface {
	IssueAgentToken(id auth.AgentIdentity, ttl time.Duration) (string, error)
}

AgentTokenMinter issues a task-scoped agent JWT for a resolved identity. It is satisfied by *auth.JWTAuthenticator (IssueAgentToken) — the same minter used at dispatch and heartbeat renewal, so the exchanged token is indistinguishable from a dispatched one on every downstream path.

type AgentTokenRenewer 

type AgentTokenRenewer interface {
	RenewAgentToken(token string, ttl, maxLifetime time.Duration) (string, bool, error)
}

AgentTokenRenewer re-mints a live attempt's agent token with a fresh short TTL, preserving the identity and the attempt's first-dispatch origin. It is consulted only on a liveness-proven heartbeat (ADR 0055 Fix #4). ok is false when the attempt has outlived its max-lifetime ceiling — the signal to let the credential lapse rather than refresh it. Implemented by *auth.JWTAuthenticator.

type Authenticator 

type Authenticator interface {
	AuthenticateAgent(token string) (*auth.AgentIdentity, error)
}

Authenticator verifies an agent bearer token into a task instance identity.

type DecisionRecorder added in v0.5.5 

type DecisionRecorder interface {
	RecordSchedulerDecision(decisionType string)
}

DecisionRecorder meters a control-plane decision by type. The production implementation is observability.Metrics (dexaflow_scheduler_decisions_total).

type InflightHandlers 

type InflightHandlers struct {
	// contains filtered or unexported fields
}

InflightHandlers counts the agent RPC handlers currently executing. gRPC exposes no such number, and the bounded graceful stop needs it: past its bound the process force-closes the transports and waits only briefly for the handlers to return, while a log writer's final Put is bounded far longer, so the process can legitimately exit with handlers still running. Abandoning one is safe — a log object is written by a single atomic Put, so the stored object stays at its previous flush rather than being truncated — but silent, and this count is what makes it visible.

The zero value is not usable; build one with NewInflightHandlers. It is safe for concurrent use.

func NewInflightHandlers 

func NewInflightHandlers() *InflightHandlers

NewInflightHandlers returns a zeroed counter to install on a gRPC server via UnaryInterceptor and StreamInterceptor and to read with Count.

func (*InflightHandlers) Count 

func (c *InflightHandlers) Count() int

Count reports how many handlers are executing right now.

func (*InflightHandlers) StreamInterceptor 

func (c *InflightHandlers) StreamInterceptor() grpc.StreamServerInterceptor

StreamInterceptor counts a streaming handler for the duration of its stream. These are the handlers the shutdown log is about: StreamLogs lives as long as its task and an idle warm worker's AwaitAssignment lives indefinitely, so they are what is still open when the bounded stop gives up.

func (*InflightHandlers) UnaryInterceptor 

func (c *InflightHandlers) UnaryInterceptor() grpc.UnaryServerInterceptor

UnaryInterceptor counts a unary handler for the duration of its call.

type LegacyTokenRecorder added in v0.5.1 

type LegacyTokenRecorder interface {
	RecordLegacyAttemptToken()
}

LegacyTokenRecorder counts agent RPCs authenticated by a task token that carries no attempt_epoch claim (ADR 0051 amendment). *observability.Metrics satisfies it.

type LogPublisher 

type LogPublisher interface {
	Publish(ctx context.Context, ref logs.Ref, line string) error
}

LogPublisher fans a log line out for live tailing (optional).

type LogResumer added in v0.5.5 

type LogResumer interface {
	Resume(ref logs.Ref) (logs.LogWriter, error)
}

LogResumer is implemented by log sinks that can continue an attempt's stored log (logs.DiskSink, logs.ObjectSink): Resume opens a writer that keeps what is already stored, where Open may start the attempt over.

type LogSink 

type LogSink interface {
	Open(ref logs.Ref) (logs.LogWriter, error)
}

LogSink opens a writer for a task attempt's streamed logs.

type LogSubscriberProbe added in v0.5.1 

type LogSubscriberProbe interface {
	HasSubscribers(ctx context.Context, ref logs.Ref) (bool, error)
}

LogSubscriberProbe is implemented by a LogPublisher that can tell whether anyone is tailing an attempt (logs.RedisTailer counts with PUBSUB NUMSUB, logs.MemoryTailer its own subscribers). With logs.tail.publish set to on_demand, StreamLogs then publishes only while someone listens; under the default (always), or with a publisher without it, every line is published.

type PodTaskResolver 

type PodTaskResolver interface {
	ResolveTaskInstance(ctx context.Context, pod ReviewedPod) (auth.AgentIdentity, error)
	ResolveAgent(ctx context.Context, pod ReviewedPod) (auth.AgentIdentity, error)
}

PodTaskResolver maps a reviewed pod to the agent identity it runs, so the minted JWT is scoped correctly. It is an interface so it is mocked in unit tests; the concrete resolver reads the pod the apiserver validated.

ResolveAgent is the branch point the exchange uses (ADR 0058 D1): a pod carrying the warm-worker LABEL resolves to a warm-worker identity (Scope == ScopeWarmWorker, naming its dag_version pool + worker id, no task claims); any other pod resolves to the task-instance identity the control plane stamped on it at dispatch — exactly what ResolveTaskInstance returns. ResolveTaskInstance is kept for the pure task path and callers that only ever expect a task instance.

type ReclaimEvent 

type ReclaimEvent struct {
	AssignmentID string
	DagVersionID string
	Reason       ReclaimReason
	RunID        string
	TaskID       string
	TryNumber    int
	// AttemptEpoch is the execution of TryNumber the assignment was dispatched as
	// (ADR 0051 amendment), so the re-place is fenced to exactly that attempt.
	AttemptEpoch int
}

ReclaimEvent is emitted when a handed-out assignment must be re-placed. The placement layer consumes it to re-dispatch the attempt on the infra budget. It carries the attempt identity (RunID, TaskID, TryNumber) — populated from the leaseState at every emit site — so the observer can re-place the exact attempt (ADR 0058 N1d-c, H2) without re-deriving it. It deliberately carries no attempt_token.

type ReclaimReason 

type ReclaimReason int

ReclaimReason names why an assignment was reclaimed (ADR 0058 N1b, H1). 

const (
	// ReclaimLeaseExpired means the worker did not ack the assignment as started
	// before its lease elapsed.
	ReclaimLeaseExpired ReclaimReason = iota
	// ReclaimRefused means the worker acked with started=false (it cannot or will
	// not run it).
	ReclaimRefused
	// ReclaimWorkerGone means the worker's stream ended while it held an unacked
	// assignment (a disconnected worker can never ack it).
	ReclaimWorkerGone
)

type ReviewedPod 

type ReviewedPod struct {
	Namespace      string
	PodName        string
	PodUID         string
	ServiceAccount string
}

ReviewedPod is the pod a validated projected ServiceAccount token identifies. The concrete TokenReviewer fills it from the Kubernetes TokenReview response — the bound-token status carries the pod name/uid the token was issued for. The field names the control plane resolves a task instance from are apiserver-version-dependent (ADR 0055 D7): the primary key is PodName in the task namespace; PodUID guards against a name-reused stale pod.

type SecretLivenessAuditor 

type SecretLivenessAuditor interface {
	// tenantID is the tenant UUID the agent token carries (AgentIdentity.TenantID),
	// not the tenant name.
	RecordSecretLivenessDenial(ctx context.Context, tenantID, dagID, runID, taskID string, tryNumber int, kind, mode string) error
}

SecretLivenessAuditor records a structured audit event when the secret-path liveness gate fires: a would-have-denied in observe mode, or a denial in enforce mode (ADR 0055). It carries identity + kind + mode only, never secret names or values. Optional and best-effort: a nil auditor or a write error only skips the row; it never changes the gate's decision.

type SecretScopeAuditor 

type SecretScopeAuditor interface {
	// tenantID is the tenant UUID the agent token carries (AgentIdentity.TenantID),
	// not the tenant name.
	RecordSecretScopeWarning(ctx context.Context, tenantID, dagID, runID, taskID, kind string, declared, total int) error
}

SecretScopeAuditor records a structured audit event when a task receives the full tenant secret set despite declaring only a subset of it — the visibility half of the warn-before-enforce arc (ADR 0045 §Settled #3, ADR 0055). It carries counts only, never secret names or values. It is optional and best-effort: a nil auditor or a write error only skips the audit row; it never changes what is delivered.

type SecretsStore 

type SecretsStore interface {
	SecretVariables(ctx context.Context, tenant string) (map[string]string, error)
	SecretConnectionURIs(ctx context.Context, tenant string) (map[string]string, error)
	SecretVariablesScoped(ctx context.Context, tenant string, names []string) (map[string]string, error)
	SecretConnectionURIsScoped(ctx context.Context, tenant string, names []string) (map[string]string, error)
}

SecretsStore returns a tenant's Variables and Connections for delivery to a task pod (ADR 0021). Connection URIs carry decrypted credentials, so this is only ever served over the authenticated agent channel — never to the UI/API.

The unscoped methods return the whole tenant vault (permissive / off). The scoped methods return ONLY the named subset, filtered IN THE QUERY (ADR 0055 D1: never post-filter the decrypted whole vault in the handler); they back secret_scoping: enforce, where a task receives only what it declared. An empty name set returns nothing — enforce's load-bearing [] case.

type Server 

type Server struct {
	agentv1.UnimplementedAgentServiceServer
	// contains filtered or unexported fields
}

Server implements agentv1.AgentServiceServer over a Store and Authenticator.

func NewServer 

func NewServer(authn Authenticator, store Store, xcomSvc XComService) *Server

NewServer builds an AgentService server backed by the given authenticator, store, and XCom service.

func (*Server) AwaitAssignment 

func (s *Server) AwaitAssignment(stream agentv1.AgentService_AwaitAssignmentServer) error

AwaitAssignment is the warm-worker assignment transport (ADR 0058 N1b): a long-lived bidi stream over which the control plane pushes per-attempt WorkAssignments down and the worker sends its registration, acks, and slot-free signals up.

Inert-when-off: if warm pools are not wired the handler refuses immediately with FailedPrecondition — the flag-gated dormant state.

Identity: the registry key is the worker's AUTHENTICATED identity from the stream's bearer token (via identify), NOT the dag_version_id in the register payload — a worker cannot claim an arbitrary identity through the message. The payload's dag_version_id only names which pool the worker serves and must be non-empty.

After registration two flows run concurrently: the receive loop drains WorkerMessages (acks feed the H1 lease machine, slot-free frees the worker) while the main select pumps assignments from the worker's outbound channel down the stream. The handler exits — deregistering the worker (defer) — on context cancellation, the control plane's shutdown signal (SetShutdown), a stream Send error, the receive loop ending (clean EOF or a transport error), a newer registration superseding this stream after it went silent past the liveness grace, or this replica losing scheduler leadership while the worker is idle (FailedPrecondition, so the worker re-dials toward the new leader). A second stream for an identity whose first stream is still live is refused with AlreadyExists (one live registration per identity).

func (*Server) EnableWarmPools 

func (s *Server) EnableWarmPools(onReclaim func(ReclaimEvent))

EnableWarmPools turns on the warm-worker assignment transport with a production registry (ADR 0058 N1b). onReclaim (may be nil) observes reclaim events for the future placement layer to consume. Call only when execution.warm_pools_enabled is set — the default leaves the handler inert.

func (*Server) ExchangeToken 

func (s *Server) ExchangeToken(ctx context.Context, _ *agentv1.ExchangeTokenRequest) (*agentv1.ExchangeTokenResponse, error)

ExchangeToken validates the agent's projected ServiceAccount token (bootstrap bearer), resolves the pod to its task instance, and returns a freshly minted task-scoped agent JWT (ADR 0055 Fix #3). It is called ONCE at agent startup under the exchange transport; the default env-var transport never calls it.

It fails closed at every step: Unimplemented when the exchange is not wired, PermissionDenied on an insecure channel, Unauthenticated on a missing or rejected projected token, and Internal when the reviewed pod cannot be resolved to an attempt (never mint an unscoped or misattributed token). The minted token and the presented projected token are never logged.

func (*Server) FetchXCom 

func (s *Server) FetchXCom(ctx context.Context, req *agentv1.FetchXComRequest) (*agentv1.FetchXComResponse, error)

FetchXCom returns an upstream task's value, but only from a task the caller declared as an XCom input within the same run (and, by construction, the same tenant), enforcing cross-tenant and cross-run isolation.

func (*Server) FetchXComBatch added in v0.5.1 

func (s *Server) FetchXComBatch(ctx context.Context, req *agentv1.FetchXComBatchRequest) (*agentv1.FetchXComBatchResponse, error)

FetchXComBatch returns several upstream values in one call. Every item is authorized like FetchXCom before anything is read: an upstream the task did not declare as an input or dependency denies the whole batch. Keys are built from the caller's token identity, so the batch is scoped to its own tenant, DAG and run exactly like FetchXCom.

func (*Server) GetConnections 

func (s *Server) GetConnections(ctx context.Context, _ *agentv1.GetConnectionsRequest) (*agentv1.GetConnectionsResponse, error)

GetConnections returns the calling task's tenant Connections as Airflow URIs for the agent to export as AIRFLOW_CONN_<CONN_ID>.

func (*Server) GetTaskSpec 

func (s *Server) GetTaskSpec(ctx context.Context, _ *agentv1.GetTaskSpecRequest) (*agentv1.TaskSpec, error)

GetTaskSpec returns the execution spec for the calling task instance.

func (*Server) GetVariables 

func (s *Server) GetVariables(ctx context.Context, _ *agentv1.GetVariablesRequest) (*agentv1.GetVariablesResponse, error)

GetVariables returns the calling task's tenant Variables for the agent to export as AIRFLOW_VAR_<KEY>.

func (*Server) Heartbeat 

func (s *Server) Heartbeat(ctx context.Context, _ *agentv1.HeartbeatRequest) (*agentv1.HeartbeatResponse, error)

Heartbeat stamps the per-TI liveness signal (#128) and returns the server clock so the agent can detect skew. A storage error stamping the heartbeat is logged but does not fail the RPC — failing the call would risk the agent terminating itself unnecessarily on a transient DB blip. The scheduler reaper would, in the worst case, fail the TI as agent_lost on the next tick; correct under "do no harm" (ADR 0031).

func (*Server) PushXCom 

func (s *Server) PushXCom(ctx context.Context, req *agentv1.PushXComRequest) (*agentv1.PushXComResponse, error)

PushXCom stores a value the task produced, keyed by the caller's identity. Size/schema violations are returned as a rejection, not a transport error, so the agent can fail the task with a clear reason.

func (*Server) Register 

func (s *Server) Register(ctx context.Context, _ *agentv1.RegisterRequest) (*agentv1.RegisterResponse, error)

Register acknowledges an agent's startup and returns the server clock.

func (*Server) ReportState 

func (s *Server) ReportState(ctx context.Context, req *agentv1.ReportStateRequest) (*agentv1.ReportStateResponse, error)

ReportState records a state transition the agent observed for its task.

func (*Server) SetLeaderCheck 

func (s *Server) SetLeaderCheck(fn func() bool)

SetLeaderCheck gates AwaitAssignment to the scheduler leader (warm-pool Hole B). Each scheduler replica wires its OWN leadership predicate, so a follower refuses the stream (FailedPrecondition) and only the leader (whose leader-only placer consults the same in-memory registry the worker registers into) serves it. An open stream re-checks the predicate, so a leader that steps down sends its idle workers away with the same code (#1602). A nil predicate (the default) leaves the handler unchecked, so a single-node or unwired deployment serves exactly as before.

func (*Server) SetLegacyTokenRecorder added in v0.5.1 

func (s *Server) SetLegacyTokenRecorder(r LegacyTokenRecorder)

SetLegacyTokenRecorder attaches the legacy-token meter. A nil recorder (the default) counts nothing; acceptance is unchanged either way. It is safe to call while the server is already serving: the boot wiring attaches it after the gRPC listener has started.

func (*Server) SetLivenessGate 

func (s *Server) SetLivenessGate(checker TaskLivenessChecker, mode string)

SetLivenessGate attaches the read-only task-instance liveness predicate the secret path consults, in the given mode ("observe" | "enforce", ADR 0055 E2). An unrecognized mode falls back to observe — the safe, non-denying default. A nil checker leaves the gate off (delivery unchanged), so the gate is opt-in.

func (*Server) SetLogPublisher 

func (s *Server) SetLogPublisher(p LogPublisher)

SetLogPublisher attaches the live-tail publisher (optional). When set, StreamLogs publishes each line for the UI's live tail.

func (*Server) SetLogSink 

func (s *Server) SetLogSink(sink LogSink)

SetLogSink attaches the log sink that StreamLogs writes to. Without it, StreamLogs reports Unimplemented.

func (*Server) SetSecretLivenessAuditor 

func (s *Server) SetSecretLivenessAuditor(a SecretLivenessAuditor)

SetSecretLivenessAuditor attaches the sink for secret-path liveness events (optional). Without it, a would-have-denied / denial still produces the WARN log but no audit row.

func (*Server) SetSecretScopeAuditor 

func (s *Server) SetSecretScopeAuditor(a SecretScopeAuditor)

SetSecretScopeAuditor attaches the sink for secret-scope warning events (optional). Without it, a narrowing declaration still produces the WARN log but no audit row.

func (*Server) SetSecretScoping 

func (s *Server) SetSecretScoping(policy string)

SetSecretScoping sets the operator scope-by-declaration policy (ADR 0055 D9): "enforce" | "permissive" | "off". An unrecognized value falls back to permissive — the safe, non-denying default — so a misconfiguration never silently denies. The policy is operator-scoped, never author-settable.

func (*Server) SetSecrets 

func (s *Server) SetSecrets(store SecretsStore, allowInsecure bool)

SetSecrets attaches the secrets store. allowInsecure permits serving secrets over a non-TLS channel — for local/dev only; production must use TLS (the handlers fail closed otherwise). See ADR 0021 / issue #58.

func (*Server) SetSettledHook added in v0.5.4 

func (s *Server) SetSettledHook(fn func())

SetSettledHook registers fn to run after an agent's terminal state report is recorded. The scheduler wires its coalescing Wake here when scheduler.eager_promotion is on; fn must return without blocking.

func (*Server) SetShutdown 

func (s *Server) SetShutdown(ctx context.Context)

SetShutdown wires the control plane's shutdown context: once ctx ends, every long-lived stream returns Unavailable, so the gRPC graceful stop that follows completes instead of waiting for the tasks themselves to finish. An open StreamLogs first closes (flushes) its log writer; the agent treats the closed stream as best-effort log delivery and keeps running its task. A StreamLogs that arrives AFTER the signal is refused before its writer is opened — the listener keeps accepting for the rest of the shutdown, and an opened-then- abandoned writer Puts an empty object over an attempt that is logging elsewhere. An open AwaitAssignment ends too — an idle warm worker holds one open indefinitely, so leaving it out would make the forced stop the normal shutdown path.

func (*Server) SetTailPublishOnDemand added in v0.5.1 

func (s *Server) SetTailPublishOnDemand(on bool)

SetTailPublishOnDemand selects logs.tail.publish: false (the default, "always") publishes every line as it arrives; true ("on_demand") publishes only while a probe finds a follower, replaying the lines held since the last probe that found none (see tailGate).

func (*Server) SetTailReplayBudget added in v0.5.8 

func (s *Server) SetTailReplayBudget(maxBytes int64, meter TailReplayMeter)

SetTailReplayBudget bounds what the on_demand gates of all streams hold for replay together (logs.tail.replay_budget_bytes, see tailBudget); 0 bounds nothing. meter, when set, receives the bytes held and the rings dropped.

func (*Server) SetTokenExchange 

func (s *Server) SetTokenExchange(reviewer TokenReviewer, resolver PodTaskResolver, minter AgentTokenMinter, ttl time.Duration, allowInsecure bool)

SetTokenExchange wires the projected-SA-token exchange (ADR 0055 Fix #3): the TokenReview client, the pod→task-instance resolver, the JWT minter, and the TTL of the minted task-scoped token. allowInsecure permits running the exchange over a non-TLS channel (dev only); production must use TLS (ExchangeToken fails closed otherwise, like the secret path). A nil reviewer leaves the exchange OFF — ExchangeToken then reports Unimplemented — which is the default (env-var) transport, so a deployment that does not opt in is byte-identical to today.

func (*Server) SetTokenRejectionRecorder added in v0.5.5 

func (s *Server) SetTokenRejectionRecorder(r TokenRejectionRecorder)

SetTokenRejectionRecorder attaches the rejection meter. A nil recorder (the default) counts nothing; the log line is written either way. It is safe to call while the server is already serving.

func (*Server) SetTokenRenewal 

func (s *Server) SetTokenRenewal(renewer AgentTokenRenewer, renewalTTL, maxAttemptLifetime time.Duration)

SetTokenRenewal wires per-attempt token renewal (ADR 0055 Fix #4): on a liveness-proven heartbeat the server re-mints the caller's bearer with a fresh renewalTTL and returns it on HeartbeatResponse.renewed_token, so a long task keeps a working credential while the short TTL bounds a stolen/finished one. maxAttemptLifetime is the hard ceiling on an attempt's total credential age since dispatch (0 disables it). A nil renewer or non-positive renewalTTL leaves renewal off — the heartbeat returns no token (unchanged behavior).

func (*Server) SetWarmPools 

func (s *Server) SetWarmPools(reg *WorkerRegistry)

SetWarmPools wires a prebuilt warm-worker assignment registry (ADR 0058 N1b). A nil registry (the default) leaves AwaitAssignment inert — it returns FailedPrecondition — so with execution.warm_pools_enabled off the transport is completely dormant and no running path changes. Used by tests to inject a registry with a deterministic lease; callers wire it via EnableWarmPools.

func (*Server) StreamLogs 

func (s *Server) StreamLogs(stream agentv1.AgentService_StreamLogsServer) (err error)

StreamLogs receives the task's log lines and writes them through the sink, flushing on stream end so the logs survive the pod. The stream also ends when the control plane shuts down (SetShutdown), so the flush runs before exit — and a stream that ARRIVES after that signal is refused before any writer is opened, so it leaves no empty object behind.

type Store 

type Store interface {
	// TaskSpec returns the execution spec for the identified task instance.
	TaskSpec(ctx context.Context, id auth.AgentIdentity) (TaskSpec, error)
	// ReportState records a state transition reported by the agent.
	ReportState(ctx context.Context, id auth.AgentIdentity, state domain.TaskState, exitCode int, errMsg string) error
	// Reschedule parks an active TI in up_for_reschedule with its next-poke time so
	// the scheduler re-dispatches it later without consuming retry budget (#380).
	Reschedule(ctx context.Context, id auth.AgentIdentity, at time.Time) error
	// RecordHeartbeat stamps last_heartbeat_at on the identified TI so the
	// scheduler's heartbeat reaper (#128) can tell live tasks from agent-lost
	// ones. The state guard inside the SQL skips already-terminal rows.
	RecordHeartbeat(ctx context.Context, id auth.AgentIdentity) error
	// BindWarmAttempt records the durable warm-attempt binding (ADR 0058 N1d-a1):
	// the warm worker pod (workerPod) that acked this attempt (runID, taskID,
	// tryNumber) as started. The write is guarded on active state, so a settled
	// attempt is never bound (a benign no-op, not an error). Called only on a warm
	// ack — with warm pools off no assignment is ever acked, so it is never called
	// and warm_worker_id stays NULL.
	BindWarmAttempt(ctx context.Context, runID, taskID string, tryNumber, attemptEpoch int, workerPod string) error
}

Store is the server's view of persistent task state.

type SweepPolicy added in v0.5.5 

type SweepPolicy int

SweepPolicy is how the registry treats a warm worker whose register does not confirm attempt_home_sweep (#1678). 

const (
	// SweepIgnore accepts every worker without looking (the read-only root is
	// off, so no sweep is expected).
	SweepIgnore SweepPolicy = iota
	// SweepWarn accepts the worker, logs a warning and meters it.
	SweepWarn
	// SweepRefuse refuses the registration, logs and meters it, so the worker
	// never joins the pool and its DAG version's attempts run on dedicated pods.
	SweepRefuse
)

type TailReplayMeter added in v0.5.8 

type TailReplayMeter interface {
	// SetLogTailReplayBytes reports the bytes every gate holds together.
	SetLogTailReplayBytes(n int64)
	// RecordLogTailReplayEviction counts a ring emptied by the budget.
	RecordLogTailReplayEviction()
}

TailReplayMeter records what the on_demand live-tail gates hold for replay (observability.Metrics implements it). Optional.

type TaskLivenessChecker 

type TaskLivenessChecker interface {
	IsTaskInstanceLive(ctx context.Context, id auth.AgentIdentity) (bool, error)
}

TaskLivenessChecker reports whether a task-instance attempt is still live — present and in an active (non-terminal) state for the given (run, task, try). It is the read-only revocation signal the secret path consults (ADR 0055 D3): a terminal, superseded, or reaped attempt is not live, so its token stops resolving secrets. The predicate derives ONLY from (run, task, try) + active state — never run recency — so a clear-and-rerun of an old run stays live.

The attempt is the caller's whole identity: (run, task, try) and its attempt epoch (ADR 0051 amendment), so an infra-replaced attempt that shares its replacement's try is not live either.

type TaskSpec 

type TaskSpec struct {
	Operator         string
	Entrypoint       string
	DagVersion       string
	Environment      map[string]string
	XComInputMapping map[string][]string
	XComSchema       map[string]any
	TimeoutSeconds   int
	// CallArgsJSON carries TaskFlow literal call args captured by the parser
	// (#115). The agent injects this verbatim as LEOFLOW_CALL_ARGS_JSON; the
	// runtime decodes it. Empty when the task has no literals. The name
	// keeps Airflow's DAG-run `params` term free for a future feature (#148).
	CallArgsJSON string
	// OperatorClass is the dotted Airflow operator/sensor class for an
	// airflow_operator task (ADR 0040); empty for native operators. The agent
	// passes it to BuildCommand to dispatch the runtime's --operator mode.
	OperatorClass string
	// OperatorArgsJSON carries the operator's constructor kwargs (JSON). The agent
	// injects it as LEOFLOW_OPERATOR_ARGS; the runtime decodes it. Empty when the
	// operator takes no args.
	OperatorArgsJSON string
	// LogicalDate is the DagRun's logical date in RFC3339; the agent derives the
	// runtime's LEOFLOW_TS/LEOFLOW_DS from it (ADR 0040). Empty leaves them unset.
	LogicalDate string
	// DependsOn lists the task's upstream task_ids. The agent fetches each one's
	// return_value so a captured operator's ti.xcom_pull(<id>) resolves it (ADR 0040).
	DependsOn []string
	// DataIntervalStart/End are the DagRun's data interval in RFC3339; the agent
	// stamps the runtime's data_interval_start/end context from them (ADR 0040).
	DataIntervalStart string
	DataIntervalEnd   string
	// ParamsJSON is the task's params, JSON-encoded (#148): the DAG's param
	// defaults, then the task's own params, then the run conf (#1841, #1844). The
	// agent stamps LEOFLOW_PARAMS so the runtime exposes context['params'] /
	// {{ params.X }}.
	ParamsJSON string
	// RunConfJSON is the DagRun's conf alone, JSON-encoded, "{}" for a run
	// without one (#1843). Unlike ParamsJSON it carries no param defaults. The
	// agent stamps LEOFLOW_RUN_CONF so the runtime exposes dag_run.conf.
	RunConfJSON string
	// RunType is the DagRun's type in the Task SDK's DagRunType vocabulary
	// (#1843). The agent stamps LEOFLOW_RUN_TYPE for dag_run.run_type.
	RunType string
	// FirstRescheduleAt is when a reschedule-mode sensor first entered reschedule,
	// RFC3339 (#380). The agent stamps LEOFLOW_FIRST_RESCHEDULE_AT so the sensor's
	// get_first_reschedule_date returns it and cumulative timeout works. Empty on
	// the first attempt (not yet rescheduled).
	FirstRescheduleAt string
	// MaxTries is the task's total attempt budget (retries + 1). The agent stamps
	// LEOFLOW_MAX_TRIES so the runtime fires on_failure_callback only on the
	// terminal attempt (#424). Zero is treated as 1 (no retries).
	MaxTries int
	// OnFailureCallback marks that the task declares an Airflow on_failure_callback
	// (#424). The agent stamps LEOFLOW_ON_FAILURE_CALLBACK=1 so the runtime runs it
	// in-process on the task's final failure.
	OnFailureCallback bool
	// DeclaredVariables and DeclaredConnections are the secret names this task
	// declared (ADR 0045, ADR 0055): the task's own set when it narrows, otherwise
	// the DAG's. They are carried on the resolved spec so a later increment can
	// scope secret delivery server-side to the declared set. Today this is data
	// only: the secret RPCs still return the whole tenant vault, so a declaration
	// changes nothing about what the agent receives.
	DeclaredVariables   []string
	DeclaredConnections []string
}

TaskSpec is the execution specification the agent needs to run a task.

type TokenRejectionRecorder added in v0.5.5 

type TokenRejectionRecorder interface {
	RecordAgentTokenRejection(reason string)
}

TokenRejectionRecorder counts agent RPCs refused for their bearer token, by reason (one of the auth.AgentToken* constants, or "missing"). *observability.Metrics satisfies it.

type TokenReviewer 

type TokenReviewer interface {
	ReviewProjectedToken(ctx context.Context, token string) (ReviewedPod, error)
}

TokenReviewer validates a projected ServiceAccount token against the control plane's audience via the Kubernetes TokenReview API and returns the pod it was issued for. It is the ONE apiserver call in the exchange, made once per pod at bootstrap (never on the secret hot path). It is an interface so it is MOCKED in unit tests — the concrete client needs a real apiserver and is exercised only by the owed real-cluster e2e. An error (bad signature, expired, wrong audience, or authenticated=false) means the token is not a valid bootstrap credential.

type WarmBinding 

type WarmBinding struct {
	RunID     string
	TaskID    string
	TryNumber int
	// AttemptEpoch is the execution of TryNumber this binding is for (ADR 0051
	// amendment), so an ack never binds a worker onto a later attempt.
	AttemptEpoch int
	PodName      string
}

WarmBinding is the durable binding an ack (started=true) establishes: the warm worker pod (PodName) now serving a specific attempt (RunID, TaskID, TryNumber). The handler persists it (BindWarmAttempt) so a later failover reaper can match bound attempts against the live warm-pod set (ADR 0058 N1d-a1). PodName is the worker's own downward-API pod name it sent in WorkerRegister — the reaper's join key against ListWarmPods — NOT the registry's authenticated identity.

type WorkerInfo added in v0.5.7 

type WorkerInfo struct {
	// PodName is WorkerRegister.pod_name.
	PodName string
	// SourceDelivery is WorkerRegister.source_delivery (ADR 0095).
	SourceDelivery bool
}

WorkerInfo is what a warm worker's registration says about it besides its identity and pool: the locator a started attempt is bound to, and what it can run.

type WorkerRegistry 

type WorkerRegistry struct {
	// contains filtered or unexported fields
}

WorkerRegistry is the concurrency-safe home of the warm-worker fleet and the H1 ack/lease machine (ADR 0058 N1b).

Data structures and complexity:

  • workers: identity -> entry. Register/Deregister/MarkFree are O(1).
  • free: dag_version -> (identity -> entry). Assign grabs one free worker of a dag_version in O(1) (a single map-range that breaks on the first element), then removes it in O(1).
  • leases: assignment_id -> in-flight lease. Ack and lease-expiry are O(1).

func NewWorkerRegistry 

func NewWorkerRegistry(onReclaim func(ReclaimEvent)) *WorkerRegistry

NewWorkerRegistry builds a registry whose reclaim events are delivered to onReclaim (may be nil).

func (*WorkerRegistry) Ack 

func (r *WorkerRegistry) Ack(assignmentID string, started bool) (*WarmBinding, bool)

Ack settles the lease for assignmentID. started=true marks the worker busy, cancels the lease (no reclaim), and returns the WarmBinding to persist — the acked attempt's (run, task, try) plus the worker's pod name — with ok=true. started=false reclaims the assignment and returns ok=false. An ack for an unknown assignment (already settled, or expired past its tombstone) also returns ok=false. Only ok=true carries a non-nil binding the handler should persist.

A late ack, one that arrives after the lease expired but while its tombstone is kept, is settled the same way: the worker runs the attempt whatever the server answers, so a started late ack still binds the attempt and marks the worker busy. A refused late ack returns the worker to the free set, which the expiry deliberately did not do.

func (*WorkerRegistry) Assign 

func (r *WorkerRegistry) Assign(dagVersion string, a *agentv1.WorkAssignment) bool

Assign hands a WorkAssignment to some free worker of dagVersion by pushing it onto that worker's outbound channel and starting its lease. It returns false when no free worker of that dag_version exists (nothing was handed out and no lease was started).

func (*WorkerRegistry) Deregister 

func (r *WorkerRegistry) Deregister(w *registeredWorker)

Deregister removes a worker's entry, but only if the registry still points at exactly this entry — a reconnect under the same identity installs a new entry, and the stale stream's later Deregister must not evict the live one. Any in-flight leases the worker still held are reclaimed: a gone worker can never ack them.

func (*WorkerRegistry) MarkFree 

func (r *WorkerRegistry) MarkFree(identity string)

MarkFree records a worker's SlotFree signal: it clears busy and returns the worker to the free set so it can take new work. Unknown identities are ignored.

func (*WorkerRegistry) Register 

func (r *WorkerRegistry) Register(identity, dagVersion, podName string, send chan *agentv1.WorkAssignment, streamDone <-chan struct{}) (*registeredWorker, error)

Register records a warm worker under its authenticated identity, ready to take work for dagVersion. podName is the worker's own pod name, carried so a started ack can bind the attempt to it. streamDone is closed when the registering stream ends (nil when the caller has no such signal). It returns the entry so the caller can Deregister exactly the entry it created.

One live registration per identity: while an earlier registration's stream is still connected and has heartbeated within workerLivenessGrace, a second Register is refused with ErrWorkerAlreadyRegistered and the live entry is left untouched. A reconnect is accepted once the earlier stream has ended (its streamDone is closed, or it already deregistered), or once it has been silent for longer than the grace; in the stale case the old entry's superseded channel is closed so its handler ends. The fresh entry starts free.

func (*WorkerRegistry) RegisterWorker added in v0.5.7 

func (r *WorkerRegistry) RegisterWorker(identity, dagVersion string, info WorkerInfo, send chan *agentv1.WorkAssignment, streamDone <-chan struct{}) (*registeredWorker, error)

RegisterWorker is Register with everything the registration states about the worker.

func (*WorkerRegistry) SetOnClaim added in v0.5.4 

func (r *WorkerRegistry) SetOnClaim(fn func())

SetOnClaim sets the hook run each time a warm worker is claimed by a started attempt (after the binding that makes it count as busy is persisted). It must not block; nil disables it. Safe to call while the registry serves workers.

func (*WorkerRegistry) SetSweepPolicy added in v0.5.5 

func (r *WorkerRegistry) SetSweepPolicy(p SweepPolicy, rec DecisionRecorder)

SetSweepPolicy sets how the registry treats a warm worker whose register does not confirm attempt_home_sweep (#1678). With execution.warm_read_only_root_filesystem on, isolation between attempts relies on the agent in the task image to empty /tmp and /dev/shm, and an agent built on an older runtime silently does not. rec (may be nil) meters each unconfirmed worker. Call it before the registry serves workers.

func (*WorkerRegistry) SweepPolicy added in v0.5.5 

func (r *WorkerRegistry) SweepPolicy() SweepPolicy

SweepPolicy reports the policy SetSweepPolicy set (SweepIgnore by default).

func (*WorkerRegistry) Touch added in v0.5.3 

func (r *WorkerRegistry) Touch(w *registeredWorker)

Touch records that w's stream is alive (any worker message, including the periodic heartbeat). It is a no-op unless the registry still points at exactly this entry, so a superseded stream cannot keep its replacement alive.

type XComService 

type XComService interface {
	Push(ctx context.Context, key xcom.Key, value []byte, contentType string, schema map[string]any) error
	Fetch(ctx context.Context, key xcom.Key) (xcom.Entry, error)
}

XComService stores and retrieves XCom values for the agent.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.