Skip to content

Commit

Permalink
Merge pull request #7 from dreadnode/feature/run-status-expansion
Browse files Browse the repository at this point in the history
Extend schemas to align with our backend changes.
  • Loading branch information
evilsocket authored Nov 19, 2024
2 parents a5ac237 + 4827e1d commit 73f489b
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 10 deletions.
4 changes: 3 additions & 1 deletion dreadnode_cli/agent/format.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,14 @@
P = t.ParamSpec("P")


def get_status_style(status: str | None) -> str:
def get_status_style(status: api.Client.StrikeRunStatus | api.Client.StrikeRunZoneStatus | None) -> str:
return (
{
"pending": "dim",
"running": "bold cyan",
"completed": "bold green",
"mixed": "bold gold3",
"terminated": "bold dark_orange3",
"failed": "bold red",
"timeout": "bold yellow",
}.get(status, "")
Expand Down
46 changes: 37 additions & 9 deletions dreadnode_cli/api.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,7 +231,25 @@ def submit_challenge_flag(self, challenge: str, flag: str) -> bool:

# Strikes

StrikeRunStatus = t.Literal["pending", "deploying", "running", "completed", "timeout", "failed"]
StrikeRunStatus = t.Literal[
"pending", # Waiting to be processed in the DB
"deploying", # Dropship pod is being created and configured
"running", # Dropship pod is actively executing
"completed", # All zones finished successfully
"mixed", # Some zones succeeded, others terminated
"terminated", # All zones ended with non-zero exit codes
"timeout", # Maximum allowed run time was exceeded
"failed", # System/infrastructure error occurred
]
StrikeRunZoneStatus = t.Literal[
"pending", # Waiting to be processed in the DB
"deploying", # Dropship is creating the zone resources
"running", # Zone pods are actively executing
"completed", # Agent completed successfully (exit code 0)
"terminated", # Agent ended with non-zero exit code
"timeout", # Maximum allowed run time was exceeded
"failed", # System/infrastructure error occurred
]

class StrikeModel(BaseModel):
key: str
Expand Down Expand Up @@ -273,7 +291,7 @@ class StrikeAgentResponse(BaseModel):
key: str
name: str | None
created_at: datetime
latest_run_status: t.Optional["Client.StrikeRunStatus"]
latest_run_status: "Client.StrikeRunStatus | None"
latest_run_id: UUID | None
versions: list["Client.StrikeAgentVersion"]
latest_version: "Client.StrikeAgentVersion"
Expand All @@ -286,7 +304,7 @@ class StrikeAgentSummaryResponse(BaseModel):
key: str
name: str | None
created_at: datetime
latest_run_status: t.Optional["Client.StrikeRunStatus"]
latest_run_status: "Client.StrikeRunStatus | None"
latest_run_id: UUID | None
latest_version: "Client.StrikeAgentVersion"
revision: int
Expand All @@ -296,23 +314,30 @@ class StrikeRunOutputScore(BaseModel):
explanation: str | None = None
metadata: dict[str, t.Any] = {}

class StrikeRunOutput(BaseModel):
data: dict[str, t.Any]
class StrikeRunOutputSummary(BaseModel):
score: t.Optional["Client.StrikeRunOutputScore"] = None
metadata: dict[str, t.Any] = {}

class StrikeRunZone(BaseModel):
class StrikeRunOutput(StrikeRunOutputSummary):
data: dict[str, t.Any]

class _StrikeRunZone(BaseModel):
id: UUID
key: str
status: "Client.StrikeRunStatus"
status: "Client.StrikeRunZoneStatus"
start: datetime | None
end: datetime | None

class StrikeRunZoneSummary(_StrikeRunZone):
outputs: list["Client.StrikeRunOutputSummary"]

class StrikeRunZone(_StrikeRunZone):
agent_logs: str | None
container_logs: dict[str, str]
outputs: list["Client.StrikeRunOutput"]
inferences: list[dict[str, t.Any]]

class StrikeRunSummaryResponse(BaseModel):
class _StrikeRun(BaseModel):
id: UUID
strike_id: UUID
strike_key: str
Expand All @@ -328,7 +353,10 @@ class StrikeRunSummaryResponse(BaseModel):
start: datetime | None
end: datetime | None

class StrikeRunResponse(StrikeRunSummaryResponse):
class StrikeRunSummaryResponse(_StrikeRun):
zones: list["Client.StrikeRunZoneSummary"]

class StrikeRunResponse(_StrikeRun):
zones: list["Client.StrikeRunZone"]

def get_strike(self, strike: str) -> StrikeResponse:
Expand Down

0 comments on commit 73f489b

Please sign in to comment.