Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Extend schemas to align with our backend changes. #7

Merged
merged 1 commit into from
Nov 19, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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
Loading